Update app.py

app.py

@@ -21,12 +21,19 @@ streamer = TextStreamer(generator.tokenizer, skip_prompt=True)
 
 def generate_response(prompt: str, max_new_tokens: int = 4096):
     try:
-        messages = [{"role": "user", "content": prompt}]
-        output = generator(messages, max_new_tokens=max_new_tokens, do_sample=False, streamer=streamer)
-        return output[0]["generated_text"][-1]["content"]
+        # Pass the prompt as a simple string, not a chat message list
+        output = generator(prompt, max_new_tokens=max_new_tokens, do_sample=False, streamer=streamer)
+        
+        # The output format is different now. We need to extract the response.
+        full_text = output[0]["generated_text"]
+        
+        # Remove the original prompt from the start of the response
+        if full_text.startswith(prompt):
+             return full_text[len(prompt):].strip()
+        return full_text
     except Exception as e:
         raise ValueError(f"Error generating response: {e}")
-
+        
 @app.post("/generate")
 async def generate_text(input: ModelInput):
     try: