File size: 2,263 Bytes
feee0c8 e3463bc feee0c8 eee75cd feee0c8 e3463bc 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 |
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()
|