Update app.py

app.py

@@ -114,13 +114,13 @@ def clean_text(text: str) -> str:
 # --- Scraping Endpoint ---
 
 @app.get("/scrape", response_model=ThreadResponse)
-def scrape(url: str = Query(...)):
+def scrape(url: str):
     scraper = cloudscraper.create_scraper()
     response = scraper.get(url)
 
     if response.status_code == 200:
-        soup = BeautifulSoup(response.content, 'html.parser')
-        comment_containers = soup.find_all('div', class_='post__content')
+        soup = BeautifulSoup(response.content, "html.parser")
+        comment_containers = soup.find_all("div", class_="post__content")
 
         if comment_containers:
             question = clean_text(comment_containers[0].get_text(strip=True, separator="\n"))
@@ -129,27 +129,46 @@ def scrape(url: str = Query(...)):
     return ThreadResponse(question="", replies=[])
 
 
-# --- Load T5-Small Model and Tokenizer ---
+# --- Load DeepSeek-R1-Distill-Qwen-1.5B Model & Tokenizer ---
 
-tokenizer = T5Tokenizer.from_pretrained("google-t5/t5-large")
-model = T5ForConditionalGeneration.from_pretrained("google-t5/t5-large")
+deepseek_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+deepseek_tokenizer = AutoTokenizer.from_pretrained(deepseek_model_name)
+deepseek_model = AutoModelForCausalLM.from_pretrained(deepseek_model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = model.to(device)
+deepseek_model = deepseek_model.to(device)
+
+
+# --- Load T5-Large Model & Tokenizer ---
 
+t5_model_name = "google-t5/t5-large"
+t5_tokenizer = T5Tokenizer.from_pretrained(t5_model_name)
+t5_model = T5ForConditionalGeneration.from_pretrained(t5_model_name)
+t5_model = t5_model.to(device)
+
+
+# --- Generation Functions ---
+
+def generate_deepseek(prompt: str) -> (str, str):
+    inputs = deepseek_tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024).to(device)
+    outputs = deepseek_model.generate(
+        **inputs,
+        max_length=512,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
+        num_return_sequences=1,
+        pad_token_id=deepseek_tokenizer.eos_token_id,
+    )
+    generated_text = deepseek_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# --- Core Generation Function Using T5 Prompting ---
+    # DeepSeek models usually do not have a special reasoning delimiter, so return empty reasoning
+    return "", generated_text.strip()
 
-def generate_text_with_t5(prompt: str) -> (str, str):
-    """
-    Accepts a prompt string that includes the T5 task prefix (e.g. "summarize: ..."),
-    generates output text, and optionally extracts reasoning if present.
-    Returns a tuple (reasoning_content, generated_text).
-    """
-    # Tokenize input prompt with truncation to max 512 tokens
-    inputs = tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
 
-    # Generate output tokens with beam search for quality
-    outputs = model.generate(
+def generate_t5(prompt: str) -> (str, str):
+    # T5 expects prompt with task prefix, e.g. "summarize: ..."
+    inputs = t5_tokenizer.encode(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
+    outputs = t5_model.generate(
         inputs,
         max_length=512,
         num_beams=4,
@@ -157,29 +176,28 @@ def generate_text_with_t5(prompt: str) -> (str, str):
         length_penalty=1.0,
         early_stopping=True,
     )
+    generated_text = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Optional: parse reasoning if your prompt/model uses a special separator like </think>
+    # Optional reasoning parsing if </think> is used
     if "</think>" in generated_text:
         reasoning_content, content = generated_text.split("</think>", 1)
-        reasoning_content = reasoning_content.strip()
-        content = content.strip()
+        return reasoning_content.strip(), content.strip()
     else:
-        reasoning_content = ""
-        content = generated_text.strip()
+        return "", generated_text.strip()
 
-    return reasoning_content, content
 
+# --- API Endpoints ---
 
-# --- /generate Endpoint Using T5 Prompting ---
+@app.post("/generate/{model_name}", response_model=GenerateResponse)
+async def generate(
+    request: PromptRequest,
+    model_name: str = Path(..., description="Model to use: 'deepseekr1-qwen' or 't5-large'")
+):
+    if model_name == "deepseekr1-qwen":
+        reasoning, text = generate_deepseek(request.prompt)
+    elif model_name == "t5-large":
+        reasoning, text = generate_t5(request.prompt)
+    else:
+        return {"reasoning_content": "", "generated_text": f"Error: Unknown model '{model_name}'."}
 
-@app.post("/generate", response_model=GenerateResponse)
-async def generate(request: PromptRequest):
-    """
-    Accepts a prompt string from frontend, which should include the T5 task prefix,
-    e.g. "summarize: {text to summarize}" or "translate English to German: {text}".
-    Returns generated text and optional reasoning content.
-    """
-    reasoning_content, generated_text = generate_text_with_t5(request.prompt)
-    return GenerateResponse(reasoning_content=reasoning_content, generated_text=generated_text)
+    return GenerateResponse(reasoning_content=reasoning, generated_text=text)