File size: 2,293 Bytes
feee0c8
 
e3463bc
feee0c8
eee75cd
feee0c8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8362905
feee0c8
 
 
 
 
 
 
e3463bc
 
8362905
feee0c8
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
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
# Gradio interface definition (without label)
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.",
  elem_id="email-input"
)


# Launch the Gradio interface
interface.launch()