Update app.py

app.py

@@ -13,24 +13,56 @@ try:
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
     print("Model and tokenizer loaded successfully!")
 
+    class Message(BaseModel):
+        role: str  # "system", "user", or "assistant"
+        content: str
+
     class GenerateRequest(BaseModel):
-        prompt: str
+        system_prompt : str
+        messages: list[Message]
         key: str
+        max_length: int = 1024
+        temperature: float = 0.7
 
     class GenerateResponse(BaseModel):
         generated_text: str
 
     @app.post("/generate", response_model=GenerateResponse)
     async def generate(request: GenerateRequest):
-        authorization = request.key
-        #token = authorization.split('Bearer ')[1]
-        if authorization != API_KEY:
+        if request.key != API_KEY:
             raise HTTPException(status_code=401, detail="Unauthorized")
+
         try:
-            output = pipe(request.prompt)[0]['generated_text']
-            return GenerateResponse(generated_text=output)
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
+            # Format messages into a prompt format the model expects
+            formatted_prompt = ""
+            formatted_prompt += f"<|system|>\n{request.system_prompt}</s>\n"
+            for message in request.messages:
+                if message.role == "system":
+                    formatted_prompt += f"<|system|>\n{message.content}</s>\n"
+                elif message.role == "user":
+                    formatted_prompt += f"<|user|>\n{message.content}</s>\n"
+                elif message.role == "assistant":
+                    formatted_prompt += f"<|assistant|>\n{message.content}</s>\n"
+            
+            # Add final assistant prefix for generation
+            formatted_prompt += "<|assistant|>\n"
+
+            output = pipe(
+                formatted_prompt,
+                max_length=request.max_length,
+                temperature=request.temperature,
+                do_sample=True
+            )[0]['generated_text']
+
+            # Extract only the newly generated assistant response
+            response_text = output.split("<|assistant|>\n")[-1].split("</s>")[0]
+        
+            return GenerateResponse(generated_text=response_text)
+        #try:
+            #output = pipe(request.prompt)[0]['generated_text']
+            #return GenerateResponse(generated_text=output)
+        #except Exception as e:
+            #raise HTTPException(status_code=500, detail=str(e))
 
 except Exception as e:
     print(f"Error: {e}")