summariser / app.py
kyserS09's picture
Update app.py
bbb8ec0 verified
raw
history blame
1.09 kB
import gradio as gr
from transformers import LEDTokenizer, LEDForConditionalGeneration
# Use Longformer Encoder-Decoder (LED) model
model_name = "allenai/led-large-16384"
tokenizer = LEDTokenizer.from_pretrained(model_name)
model = LEDForConditionalGeneration.from_pretrained(model_name)
def summarize_text(text):
# Tokenize input with truncation to fit within 16,384 tokens
inputs = tokenizer([text], max_length=16384, return_tensors="pt", truncation=True)
# Generate summary with adjusted parameters
summary_ids = model.generate(
inputs["input_ids"],
num_beams=4,
max_length=512, # Can be adjusted based on summary size needs
min_length=100,
early_stopping=True
)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
# Gradio Interface
iface = gr.Interface(
fn=summarize_text,
inputs="text",
outputs="text",
title="Longformer Summarizer",
description="Enter text to get a summary using the Longformer Encoder-Decoder."
)
if __name__ == "__main__":
iface.launch()