Update app.py

app.py

@@ -1,27 +1,20 @@
 import gradio as gr
-from transformers import T5ForConditionalGeneration, T5Tokenizer
+from transformers import BartForConditionalGeneration, BartTokenizer
 
-# Load the T5 model and tokenizer for question generation
-model_name = "valhalla/t5-small-qg-prepend"
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
+# Load the BART model and tokenizer
+model_name = "facebook/bart-large-cnn"
+tokenizer = BartTokenizer.from_pretrained(model_name)
+model = BartForConditionalGeneration.from_pretrained(model_name)
 
 def generate_questions(email_text):
-    # Prepend "generate questions: " to the input text
-    input_text = "generate questions: " + email_text
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
+    # Preprocess the email text for the BART model
+    inputs = tokenizer(email_text, return_tensors="pt", max_length=1024, truncation=True)
 
     # Generate questions
-    outputs = model.generate(
-        input_ids=input_ids,
-        max_length=512,
-        num_beams=4,
-        early_stopping=True
-    )
+    summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=50, early_stopping=True)
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
-    # Decode the generated text
-    questions = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return questions
+    return summary
 
 # Create a Gradio interface
 iface = gr.Interface(