Update app.py

app.py

@@ -28,19 +28,29 @@ from fastapi import Form
 class Query(BaseModel):
     text: str
    
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-tokenizer = BertTokenizerFast.from_pretrained('mrm8488/bert-small2bert-small-finetuned-cnn_daily_mail-summarization')
-model = EncoderDecoderModel.from_pretrained('mrm8488/bert-small2bert-small-finetuned-cnn_daily_mail-summarization').to(device)
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# tokenizer = BertTokenizerFast.from_pretrained('mrm8488/bert-small2bert-small-finetuned-cnn_daily_mail-summarization')
+# model = EncoderDecoderModel.from_pretrained('mrm8488/bert-small2bert-small-finetuned-cnn_daily_mail-summarization').to(device)
+
+summarizer = pipeline(
+    "summarization",
+    "pszemraj/long-t5-tglobal-base-16384-book-summary",
+    device=0 if torch.cuda.is_available() else -1,
+)
+
 
 def generate_summary(text):
+    
+    result = summarizer(text)
+    return result[0]["summary_text"]
     # cut off at BERT max length 512
-    inputs = tokenizer([text], padding="max_length", truncation=True, max_length=512, return_tensors="pt")
-    input_ids = inputs.input_ids.to(device)
-    attention_mask = inputs.attention_mask.to(device)
+    # inputs = tokenizer([text], padding="max_length", truncation=True, max_length=512, return_tensors="pt")
+    # input_ids = inputs.input_ids.to(device)
+    # attention_mask = inputs.attention_mask.to(device)
 
-    output = model.generate(input_ids, attention_mask=attention_mask)
+    # output = model.generate(input_ids, attention_mask=attention_mask)
 
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    # return tokenizer.decode(output[0], skip_special_tokens=True)
 
 
 from fastapi import FastAPI, Request, Depends, UploadFile, File
@@ -93,6 +103,4 @@ async def get_answer(q: Query ):
       
    
     return "hello"
-    
-
-
+