Update app.py

app.py

@@ -1,28 +1,44 @@
+import re
 import gradio as gr
-from transformers import BartForConditionalGeneration, BartTokenizer
+from transformers import pipeline, BartTokenizer, BartForConditionalGeneration
 
-# Load the BART model and tokenizer
+# Load the BART model and tokenizer for text generation (answer suggestions)
 model_name = "facebook/bart-large-cnn"
 tokenizer = BartTokenizer.from_pretrained(model_name)
 model = BartForConditionalGeneration.from_pretrained(model_name)
 
-def generate_questions(email_text):
-    # Preprocess the email text for the BART model
-    inputs = tokenizer(email_text, return_tensors="pt", max_length=1024, truncation=True)
+# Question detection function
+def detect_questions(email_text):
+    # Use regex to find questions in the email
+    questions = re.findall(r'([^\.\!\?]*\?)', email_text)
+    return questions
 
-    # Generate questions
-    summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=500, early_stopping=True)
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+# Generate answers using the BART model
+def generate_answers(question):
+    # Use the BART model to generate a response
+    inputs = tokenizer(question, return_tensors="pt", max_length=1024, truncation=True)
+    summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=50, early_stopping=True)
+    answer = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return answer
 
-    return summary
+# Main function to handle the email input
+def process_email(email_text):
+    questions = detect_questions(email_text)
+    responses = {}
+    
+    for question in questions:
+        response = generate_answers(question)
+        responses[question] = response
+    
+    return responses
 
 # Create a Gradio interface
 iface = gr.Interface(
-    fn=generate_questions,
+    fn=process_email,
     inputs="textbox",
-    outputs="textbox",
-    title="Email Question Generator",
-    description="Input an email, and the AI will generate the biggest questions that probably need to be answered.",
+    outputs="json",
+    title="Email Question Detector and Responder",
+    description="Input an email, and the AI will detect questions and provide suggested responses.",
 )
 
 # Launch the interface