text_summarise / app.py
vhr1007's picture
Update app.py
a8d0530 verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
model_name = "huggingface/llama-model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
def chunk_text(text, chunk_size=512):
tokens = tokenizer.encode(text, return_tensors="pt", truncation=False)
chunks = [tokens[0][i:i + chunk_size] for i in range(0, tokens.size(1), chunk_size)]
return chunks
def summarize_chunk(chunk, max_length=50):
summary_ids = model.generate(chunk.unsqueeze(0), max_length=max_length, min_length=25, length_penalty=2.0, num_beams=4, early_stopping=True)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
def summarize(text, max_summary_length=50):
chunks = chunk_text(text)
summaries = [summarize_chunk(chunk, max_summary_length) for chunk in chunks]
combined_summary = " ".join(summaries)
final_summary = summarize_chunk(tokenizer.encode(combined_summary, return_tensors="pt", truncation=True)[0], max_length=max_summary_length)
return final_summary
iface = gr.Interface(
fn=summarize,
inputs=[
gr.inputs.Textbox(lines=10, label="Input Text"),
gr.inputs.Slider(minimum=10, maximum=100, default=50, label="Max Summary Length (Optional)")
],
outputs="text",
title="Concise Text Summarization with Llama"
)
if __name__ == "__main__":
iface.launch()