Update app.py

app.py

@@ -1,16 +1,134 @@
-from transformers import T5ForConditionalGeneration, T5Tokenizer
+from fastapi import FastAPI, HTTPException, Header, Depends
+from pydantic import BaseModel
+from typing import Optional, List
+from datetime import datetime
+import torch
+from transformers import BartForConditionalGeneration, BartTokenizer
+import time
 
-model_name = "t5-base"  # Change model as needed
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
+app = FastAPI()
 
-def paraphrase(text, num_variations=1, style="standard"):
-    input_text = f"paraphrase: {text} </s>"
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
+# Configuration
+API_KEYS = {
+    "bdLFqk4IcYmRE2ONZeCts4DWrqkpqQxW": "user1"  # In production, use a secure database
+}
 
-    outputs = model.generate(**inputs, max_length=150, num_return_sequences=num_variations, temperature=1.5, top_k=100)
-    paraphrased_texts = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+# Initialize model and tokenizer
+MODEL_NAME = "facebook/bart-large-cnn"
+print("Loading model and tokenizer...")
+tokenizer = BartTokenizer.from_pretrained(MODEL_NAME, cache_dir="model_cache")
+model = BartForConditionalGeneration.from_pretrained(MODEL_NAME, cache_dir="model_cache")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = model.to(device)
+print(f"Model and tokenizer loaded successfully on {device}!")
 
-    return paraphrased_texts
+class TextRequest(BaseModel):
+    text: str
+    max_length: Optional[int] = 150
+    min_length: Optional[int] = 40
+    num_variations: Optional[int] = 1
 
-print(paraphrase("We are a company that uses our skills and creativity to assist businesses in expanding online."))
+class BatchRequest(BaseModel):
+    texts: List[str]
+    max_length: Optional[int] = 150
+    min_length: Optional[int] = 40
+    num_variations: Optional[int] = 1
+
+async def verify_api_key(api_key: str = Header(..., name="X-API-Key")):
+    if api_key not in API_KEYS:
+        raise HTTPException(status_code=403, detail="Invalid API key")
+    return api_key
+
+def generate_summary(text: str, max_length: int = 150, min_length: int = 40, num_variations: int = 1) -> List[str]:
+    try:
+        # Tokenize the input text
+        inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt").to(device)
+        
+        # Generate summaries
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=max_length,
+                min_length=min_length,
+                num_return_sequences=num_variations,
+                num_beams=num_variations * 2,
+                early_stopping=True,
+                diversity_penalty=0.5 if num_variations > 1 else 0.0,
+                num_beam_groups=num_variations if num_variations > 1 else 1
+            )
+        
+        # Decode the generated outputs
+        summaries = [
+            tokenizer.decode(output, skip_special_tokens=True) 
+            for output in outputs
+        ]
+        
+        return summaries
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Summary generation error: {str(e)}")
+
+@app.get("/")
+async def root():
+    return {"message": "Summarization API is running. Use /docs for API documentation."}
+
+@app.post("/api/summarize")
+async def summarize(request: TextRequest, api_key: str = Depends(verify_api_key)):
+    try:
+        start_time = time.time()
+        
+        summaries = generate_summary(
+            request.text, 
+            request.max_length, 
+            request.min_length,
+            request.num_variations
+        )
+        
+        processing_time = time.time() - start_time
+        
+        return {
+            "status": "success",
+            "original_text": request.text,
+            "summarized_texts": summaries,
+            "max_length": request.max_length,
+            "min_length": request.min_length,
+            "processing_time": f"{processing_time:.2f} seconds",
+            "timestamp": datetime.now().isoformat()
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/batch-summarize")
+async def batch_summarize(request: BatchRequest, api_key: str = Depends(verify_api_key)):
+    try:
+        start_time = time.time()
+        results = []
+        
+        for text in request.texts:
+            summaries = generate_summary(
+                text, 
+                request.max_length, 
+                request.min_length,
+                request.num_variations
+            )
+            
+            results.append({
+                "original_text": text,
+                "summarized_texts": summaries,
+                "max_length": request.max_length,
+                "min_length": request.min_length
+            })
+        
+        processing_time = time.time() - start_time
+        
+        return {
+            "status": "success",
+            "results": results,
+            "total_texts_processed": len(request.texts),
+            "processing_time": f"{processing_time:.2f} seconds",
+            "timestamp": datetime.now().isoformat()
+        }
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))