Update main.py

main.py

@@ -1,9 +1,11 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
-# Load the summarization model
-summarizer = pipeline("summarization", model="t5-small")
+# Load the model and tokenizer
+model_name = "shahzaib201/AI_OEL"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Pydantic model for input validation
 class TextInput(BaseModel):
@@ -16,5 +18,13 @@ app = FastAPI()
 # Endpoint for text summarization
 @app.post("/summarize_text")
 async def summarize_text_endpoint(item: TextInput):
-    summary = summarizer(item.text, max_length=item.max_length, min_length=30, do_sample=False)[0]['summary_text']
+    # Tokenize the input text
+    inputs = tokenizer(item.text, return_tensors="pt", max_length=1024, truncation=True)
+    
+    # Generate the summary
+    summary_ids = model.generate(inputs.input_ids, max_length=item.max_length, num_beams=4, length_penalty=2.0, early_stopping=True)
+    
+    # Decode the generated summary
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    
     return {"summary": summary}