Update app.py

app.py

@@ -83,7 +83,7 @@ from fastapi import FastAPI, Query, Path
 from pydantic import BaseModel
 import cloudscraper
 from bs4 import BeautifulSoup
-from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForCausalLM, T5Tokenizer, T5ForConditionalGeneration, PegasusTokenizer, PegasusForConditionalGeneration
 import torch
 import re
 from fastapi.responses import JSONResponse
@@ -144,6 +144,13 @@ t5_tokenizer = T5Tokenizer.from_pretrained(t5_model_name)
 t5_model = T5ForConditionalGeneration.from_pretrained(t5_model_name)
 t5_model = t5_model.to(device)
 
+pegasus_model_name = "google/pegasus-large"
+pegasus_tokenizer = PegasusTokenizer.from_pretrained(pegasus_model_name)
+pegasus_model = PegasusForConditionalGeneration.from_pretrained(pegasus_model_name)
+pegasus_model = pegasus_model.to(device)
+
+
+
 # --- Generation Functions ---
 
 def generate_deepseek(prompt: str) -> (str, str):
@@ -185,20 +192,45 @@ def generate_t5(prompt: str) -> (str, str):
 
 # --- API Endpoints ---
 
+def generate_pegasus(prompt: str) -> (str, str):
+    # Pegasus expects raw text input (no prefix needed)
+    inputs = pegasus_tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=1024,
+    ).to(device)
+
+    outputs = pegasus_model.generate(
+        **inputs,
+        max_new_tokens=150,
+        num_beams=4,
+        length_penalty=2.0,
+        early_stopping=True,
+    )
+    generated_text = pegasus_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Pegasus does not use <think> tags, so no reasoning extraction
+    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' or 't5-large'")
+    model_name: str = Path(..., description="Model to use: 'deepseekr1-qwen', 't5-large' or 'pegasus-large'")
 ):
     if model_name == "deepseekr1-qwen":
         reasoning, text = generate_deepseek(request.prompt)
     elif model_name == "t5-large":
         reasoning, text = generate_t5(request.prompt)
+    elif model_name == "pegasus-large":
+        reasoning, text = generate_pegasus(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)