Update app.py

app.py

@@ -38,12 +38,12 @@ def scrape(url: str = Query(...)):
 
 MODEL_NAME = "microsoft/phi-2"
 
-# Load the pipeline once at startup with device auto-mapping
+# Load the text-generation pipeline once at startup
 text_generator = pipeline(
     "text-generation",
     model=MODEL_NAME,
     trust_remote_code=True,
-    device=0 if torch.cuda.is_available() else -1,
+    device=0 if torch.cuda.is_available() else -1,  # GPU if available, else CPU
 )
 
 class PromptRequest(BaseModel):
@@ -51,13 +51,16 @@ class PromptRequest(BaseModel):
 
 @app.post("/generate")
 async def generate_text(request: PromptRequest):
-    # Prepare messages as expected by the model pipeline
-    messages = [{"role": "user", "content": request.prompt}]
+    # The model expects a string prompt, so pass request.prompt directly
+    outputs = text_generator(
+        request.prompt,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
+        num_return_sequences=1,
+    )
     
-    # 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']
 
     # Optional: parse reasoning and content if your model uses special tags like </think>