Update app.py

app.py

@@ -2,11 +2,8 @@ from fastapi import FastAPI, Query
 from pydantic import BaseModel
 import cloudscraper
 from bs4 import BeautifulSoup
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import pipeline
 import re
-import os
-
 
 app = FastAPI()
 
@@ -34,39 +31,34 @@ def scrape(url: str = Query(...)):
             return ThreadResponse(question=question, replies=replies)
     return ThreadResponse(question="", replies=[])
 
-
 MODEL_NAME = "google/flan-t5-small"
 
-# Load tokenizer and model once at startup, with device auto-mapping
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype="auto", device_map="auto")
-model.eval()
+# Load the pipeline once at startup with device auto-mapping
+text_generator = pipeline(
+    "text2text-generation",
+    model=MODEL_NAME,
+    device=0 if torch.cuda.is_available() else -1,
+    max_new_tokens=512,
+    temperature=0.5
+)
 
 class PromptRequest(BaseModel):
     prompt: str
 
 @app.post("/generate")
 async def generate_text(request: PromptRequest):
-    # Prepare chat-style input with thinking mode enabled
-    messages = [{"role": "user", "content": request.prompt}]
-    text = tokenizer.apply_chat_template(messages, tokenize=False, enable_thinking=True)
-
-    inputs = tokenizer([text], return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=512, temperature=0.5)
-    output_ids = generated_ids[:, inputs.input_ids.shape[-1]:].tolist()[0]
-    output_text = tokenizer.decode(output_ids)
+    # Use the pipeline to generate text directly
+    output = text_generator(request.prompt)[0]['generated_text']
 
     # Extract reasoning and content parts if thinking tags are present
-    if "</think>" in output_text:
-        reasoning_content = output_text.split("</think>")[0].strip()
-        content = output_text.split("</think>")[1].strip().rstrip("</s>")
+    if "</think>" in output:
+        reasoning_content = output.split("</think>")[0].strip()
+        content = output.split("</think>")[1].strip().rstrip("</s>")
     else:
         reasoning_content = ""
-        content = output_text.strip().rstrip("</s>")
+        content = output.strip().rstrip("</s>")
 
     return {
         "reasoning_content": reasoning_content,
         "generated_text": content
     }
-