Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import streamlit as st
 import fitz  # PyMuPDF for PDF text extraction
 
@@ -19,8 +19,14 @@ model_name = model_options[model_choice]
 
 @st.cache_resource(show_spinner=True)
 def load_model(name):
-    task = "text2text-generation" if "t5" in name.lower() or "pegasus" in name.lower() else "text-generation"
-    return pipeline(task, model=name)
+    if "t5" in name.lower():
+        tokenizer = AutoTokenizer.from_pretrained(name, use_fast=False)
+        model = AutoModelForSeq2SeqLM.from_pretrained(name)
+        return pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+    elif "pegasus" in name.lower():
+        return pipeline("text2text-generation", model=name)
+    else:
+        return pipeline("text-generation", model=name)
 
 simplifier = load_model(model_name)