FastMail / app.py
MNGames's picture
Update app.py
486c2a9 verified
raw
history blame
No virus
1 kB
import gradio as gr
from transformers import BartForConditionalGeneration, BartTokenizer
# Load the BART model and tokenizer
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)
# Generate questions
summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=50, early_stopping=True)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
# Create a Gradio interface
iface = gr.Interface(
fn=generate_questions,
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.",
)
# Launch the interface
iface.launch()