Update app.py

app.py

@@ -149,6 +149,10 @@ pegasus_tokenizer = PegasusTokenizer.from_pretrained(pegasus_model_name)
 pegasus_model = PegasusForConditionalGeneration.from_pretrained(pegasus_model_name)
 pegasus_model = pegasus_model.to(device)
 
+qwen3_model_name = "Qwen/Qwen3-0.6B"
+qwen3_tokenizer = AutoTokenizer.from_pretrained(qwen3_model_name)
+qwen3_model = AutoModelForCausalLM.from_pretrained(qwen3_model_name)
+qwen3_model = qwen3_model.to(device)
 
 
 # --- Generation Functions ---
@@ -213,11 +217,37 @@ def generate_pegasus(prompt: str) -> (str, str):
     # Pegasus does not use <think> tags, so no reasoning extraction
     return "", generated_text.strip()
 
+def generate_qwen3(prompt: str) -> (str, str):
+    inputs = qwen3_tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=1024,
+    ).to(device)
+
+    outputs = qwen3_model.generate(
+        **inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
+        num_return_sequences=1,
+        pad_token_id=qwen3_tokenizer.eos_token_id,
+    )
+
+    generated_text = qwen3_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    if "</think>" in generated_text:
+        reasoning_content, content = generated_text.split("</think>", 1)
+        return reasoning_content.strip(), content.strip()
+    else:
+        return "", generated_text.strip()
+
 
 @app.post("/generate/{model_name}", response_model=GenerateResponse)
 async def generate(
     request: PromptRequest,
-    model_name: str = Path(..., description="Model to use: 'deepseekr1-qwen', 't5-large' or 'pegasus-large'")
+    model_name: str = Path(..., description="Model to use: 'deepseekr1-qwen', 't5-large', 'pegasus-large', or 'qwen3-0.6b'")
 ):
     if model_name == "deepseekr1-qwen":
         reasoning, text = generate_deepseek(request.prompt)
@@ -225,12 +255,15 @@ async def generate(
         reasoning, text = generate_t5(request.prompt)
     elif model_name == "pegasus-large":
         reasoning, text = generate_pegasus(request.prompt)
+    elif model_name == "qwen3-0.6b":
+        reasoning, text = generate_qwen3(request.prompt)
     else:
         return GenerateResponse(reasoning_content="", generated_text=f"Error: Unknown model '{model_name}'.")
 
     return GenerateResponse(reasoning_content=reasoning, generated_text=text)
 
 
+
 # --- Global Exception Handler ---
 
 @app.exception_handler(Exception)