Update app.py

app.py

@@ -1,43 +1,64 @@
-import gradio as gr
-from transformers import pipeline, BartTokenizer, BartForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from gradio import Interface
 
-# Load the BART model and tokenizer for text generation
+# Define the model name (change if desired)
 model_name = "facebook/bart-base"
-tokenizer = BartTokenizer.from_pretrained(model_name)
-model = BartForConditionalGeneration.from_pretrained(model_name)
-
-def detect_questions(email_text):
-    # Use the BART model to generate questions from the email text
-    inputs = tokenizer("generate questions: " + email_text, return_tensors="pt", max_length=1024, truncation=True)
-    questions = model.generate(inputs["input_ids"], num_beams=4, max_length=50, early_stopping=True)
-    questions = tokenizer.decode(questions[0], skip_special_tokens=True)
-    return questions.split("##")
-
-def generate_responses(questions):
-    responses = {}
-    for question in questions:
-        # Generate a response for each question using the BART model
-        inputs = tokenizer(question, return_tensors="pt", max_length=1024, truncation=True)
-        response = model.generate(inputs["input_ids"], num_beams=4, max_length=200, early_stopping=True)
-        response = tokenizer.decode(response[0], skip_special_tokens=True)
-        responses[question] = response
-    return responses
-
-def process_email(email_text):
-    # Detect questions from the email
-    questions = detect_questions(email_text)
-    
-    # Generate responses to the detected questions
-    responses = generate_responses(questions)
-    
-    return responses
-
-iface = gr.Interface(
-    fn=process_email,
-    inputs="textbox",
-    outputs="text",
-    title="Email Question Detector and Responder",
-    description="Input an email, and the AI will detect questions and provide possible responses.",
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+def generate_questions(email):
+  """Generates questions based on the input email."""
+  # Encode the email with tokenizer
+  inputs = tokenizer(email, return_tensors="pt")
+
+  # Generate questions using model with specific prompt
+  generation = model.generate(
+      input_ids=inputs["input_ids"],
+      max_length=256,  # Adjust max length as needed
+      num_beams=5,  # Adjust beam search for better quality (slower)
+      early_stopping=True,
+      prompt="What are the important questions or things that need to be addressed in this email:\n",
+  )
+
+  # Decode the generated text
+  return tokenizer.decode(generation[0], skip_special_tokens=True)
+
+def generate_answers(questions):
+  """Generates possible answers to the input questions."""
+  # Encode each question with tokenizer, separated by newline
+  inputs = tokenizer("\n".join(questions), return_tensors="pt")
+
+  # Generate answers using model with specific prompt
+  generation = model.generate(
+      input_ids=inputs["input_ids"],
+      max_length=512,  # Adjust max length as needed
+      num_beams=3,  # Adjust beam search for better quality (slower)
+      early_stopping=True,
+      prompt="Here are some possible answers to the questions:\n",
+  )
+
+  # Decode the generated text
+  answers = tokenizer.decode(generation[0], skip_special_tokens=True).split("\n")
+  return zip(questions, answers[1:])  # Skip the first answer (prompt repetition)
+
+def gradio_app(email):
+  """Gradio interface function"""
+  questions = generate_questions(email)
+  answers = generate_answers(questions.split("\n"))
+  return questions, [answer for _, answer in answers]
+
+# Gradio interface definition
+interface = Interface(
+  fn=gradio_app,
+  inputs="textbox",
+  outputs=["text", "text"],
+  title="AI Email Assistant",
+  description="Enter a long email and get questions and possible answers generated by an AI model.",
+  label="Email",
+  elem_id="email-input"
 )
 
-iface.launch()
+# Launch the Gradio interface
+interface.launch()