mt5-small / app.py
JDmayo's picture
Update app.py
2c8d87e verified
import gradio as gr
import subprocess
import sys
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
# Ensure sentencepiece is installed
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'sentencepiece'])
# Load the tokenizer and model from the downloaded directory
model_name_or_path = 'model_directory'
try:
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
except ValueError as e:
print(f"Error loading fast tokenizer: {e}. Trying to load slow tokenizer.")
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
# Define the inference function
def generate_summary(text):
inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)
summary_ids = model.generate(inputs, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)
return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
# Define the Gradio interface
def inference(text):
summary = generate_summary(text)
return summary
interface = gr.Interface(fn=inference, inputs="text", outputs="text", title="Text Summarization", description="Enter text to summarize")
# Launch the Gradio interface
interface.launch()