JDmayo/mt5-small

1import gradio as gr2import subprocess3import sys4from transformers import AutoModelForSeq2SeqLM, AutoTokenizer5 6# Ensure sentencepiece is installed7subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'sentencepiece'])8 9# Load the tokenizer and model from the downloaded directory10model_name_or_path = 'model_directory'11 12try:13    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)14except ValueError as e:15    print(f"Error loading fast tokenizer: {e}. Trying to load slow tokenizer.")16    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=False)17 18model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)19 20# Define the inference function21def generate_summary(text):22    inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)23    summary_ids = model.generate(inputs, max_length=150, min_length=40, length_penalty=2.0, num_beams=4, early_stopping=True)24    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)25 26# Define the Gradio interface27def inference(text):28    summary = generate_summary(text)29    return summary30 31interface = gr.Interface(fn=inference, inputs="text", outputs="text", title="Text Summarization", description="Enter text to summarize")32 33# Launch the Gradio interface34interface.launch()35