from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from gradio import Interface

# Define the model name (change if desired)
model_name = "facebook/bart-base"

# 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"
)

# Launch the Gradio interface
interface.launch()