Update app.py

app.py

@@ -36,15 +36,14 @@ def scrape(url: str = Query(...)):
             return ThreadResponse(question=question, replies=replies)
     return ThreadResponse(question="", replies=[])
 
-MODEL_NAME = "google/flan-t5-small"
+MODEL_NAME = "deepseek-ai/DeepSeek-R1"
 
 # Load the pipeline once at startup with device auto-mapping
 text_generator = pipeline(
-    "text2text-generation",
+    "text-generation",
     model=MODEL_NAME,
+    trust_remote_code=True,
     device=0 if torch.cuda.is_available() else -1,
-    max_new_tokens=512,
-    temperature=0.5
 )
 
 class PromptRequest(BaseModel):
@@ -52,18 +51,24 @@ class PromptRequest(BaseModel):
 
 @app.post("/generate")
 async def generate_text(request: PromptRequest):
-    # Use the pipeline to generate text directly
-    output = text_generator(request.prompt)[0]['generated_text']
+    # Prepare messages as expected by the model pipeline
+    messages = [{"role": "user", "content": request.prompt}]
+    
+    # Call the pipeline with messages
+    outputs = text_generator(messages)
+    
+    # The pipeline returns a list of dicts with 'generated_text'
+    generated_text = outputs[0]['generated_text']
 
-    # Extract reasoning and content parts if thinking tags are present
-    if "</think>" in output:
-        reasoning_content = output.split("</think>")[0].strip()
-        content = output.split("</think>")[1].strip().rstrip("</s>")
+    # Optional: parse reasoning and content if your model uses special tags like </think>
+    if "</think>" in generated_text:
+        reasoning_content = generated_text.split("</think>")[0].strip()
+        content = generated_text.split("</think>")[1].strip()
     else:
         reasoning_content = ""
-        content = output.strip().rstrip("</s>")
+        content = generated_text.strip()
 
     return {
         "reasoning_content": reasoning_content,
         "generated_text": content
-    }
+    }