|
|
import gradio as gr |
|
|
import subprocess |
|
|
import sys |
|
|
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer |
|
|
|
|
|
|
|
|
subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'sentencepiece']) |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
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") |
|
|
|
|
|
|
|
|
interface.launch() |
|
|
|