Fixed app v2

app.py

@@ -1,44 +1,36 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import pipeline, set_seed
-from typing import List
+from transformers import pipeline
+from typing import List, Dict
 
 # Initialize the FastAPI app
 app = FastAPI()
 
-# Initialize the generator pipeline
-generator = pipeline('text-generation', model='gpt2-medium')
-set_seed(42)
+# Initialize the text generation pipeline
+pipe = pipeline("text-generation", model="CyberNative-AI/Colibri_8b_v0.1")
 
-# Data model for FastAPI input
-class UserInput(BaseModel):
-    conversation: str
-    user_input: str
-    max_length: int = 50  # default length
-    num_return_sequences: int = 1  # default number of sequences
+# Define the input schema for FastAPI
+class Message(BaseModel):
+    role: str
+    content: str
+
+class MessagesInput(BaseModel):
+    messages: List[Message]
 
 @app.post("/generate/")
-async def generate_response(user_input: UserInput):
+async def generate_response(messages_input: MessagesInput):
     try:
-        # Construct the prompt from the conversation and user input
-        prompt = f"{user_input.conversation}{user_input.user_input}"
-
-        # Generate response
-        responses = generator(
-            prompt,
-            max_length=user_input.max_length,
-            num_return_sequences=user_input.num_return_sequences
-        )
-
-        # Extract text from each generated sequence
-        generated_texts = [response["generated_text"] for response in responses]
+        # Convert messages to the expected format
+        messages = [{"role": msg.role, "content": msg.content} for msg in messages_input.messages]
+
+        # Generate response using the pipeline
+        response = pipe(messages)
         
-        # Update conversation with the last generated text
-        updated_conversation = f"{prompt}\n{generated_texts[-1]}"
+        # Extract generated text
+        generated_text = response[0]["generated_text"]
 
         return {
-            "responses": generated_texts,
-            "updated_conversation": updated_conversation
+            "response": generated_text
         }
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))