Update app.py

app.py

@@ -1,20 +1,38 @@
-from fastapi import FastAPI
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import pipeline
 import torch
 
-app = FastAPI()
+app = FastAPI(title="Text Summarization API")
 
-MODEL_NAME = "deepseek-ai/deepseek-llm-7b"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
+summarizer = pipeline(
+    "summarization",
+    model="deepseek-ai/deepseek-llm-7b-base",
+    device=0 if torch.cuda.is_available() else -1
+)
 
-@app.get("/")
-def root():
-    return {"message": "DeepSeek API is running!"}
+class SummarizationRequest(BaseModel):
+    text: str
+    max_length: int = 130
+    min_length: int = 20
 
-@app.post("/summarize/")
-async def summarize(input_text: str):
-    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
-    output = model.generate(**inputs, max_new_tokens=100)
-    summary = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"summary": summary}
+class SummarizationResponse(BaseModel):
+    summary: str
+
+@app.post("/summarize", response_model=SummarizationResponse)
+async def summarize_text(request: SummarizationRequest):
+    try:
+        summary = summarizer(
+            request.text,
+            max_length=request.max_length,
+            min_length=request.min_length,
+            do_sample=False
+        )[0]['summary_text']
+        
+        return SummarizationResponse(summary=summary)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy"}