File size: 3,791 Bytes
cee775f cadff73 cee775f dcfe9dc cee775f dcfe9dc cee775f dcfe9dc cee775f cadff73 cee775f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import torch
import gradio as gr
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer, AutoConfig
# List of summarization models
model_names = [
"google/bigbird-pegasus-large-arxiv",
"facebook/bart-large-cnn",
"google/t5-v1_1-large",
"sshleifer/distilbart-cnn-12-6",
"allenai/led-base-16384",
"google/pegasus-xsum",
"togethercomputer/LLaMA-2-7B-32K"
]
# Placeholder for the summarizer pipeline, tokenizer, and maximum tokens
summarizer = None
tokenizer = None
max_tokens = None
# Example text for summarization
example_text = (
"Artificial intelligence (AI) is intelligence—perceiving, synthesizing, and inferring information—"
"demonstrated by machines, as opposed to intelligence displayed by non-human animals and humans. "
"Example tasks in which AI is employed include speech recognition, computer vision, language translation, "
"autonomous vehicles, and game playing. AI research has been defined as the field of study of intelligent "
"agents, which refers to any system that perceives its environment and takes actions that maximize its "
"chance of achieving its goals."
)
# Function to load the selected model
def load_model(model_name):
global summarizer, tokenizer, max_tokens
try:
# Load the summarization pipeline with the selected model
summarizer = pipeline("summarization", model=model_name, torch_dtype=torch.float32)
tokenizer = AutoTokenizer.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name)
# Set a reasonable default for max_tokens if not available
max_tokens = getattr(config, 'max_position_embeddings', 1024)
return f"Model {model_name} loaded successfully! Max tokens: {max_tokens}"
except Exception as e:
return f"Failed to load model {model_name}. Error: {str(e)}"
# Function to summarize the input text
def summarize_text(input, min_length, max_length):
if summarizer is None:
return "No model loaded!"
try:
# Tokenize the input text and check the number of tokens
input_tokens = tokenizer.encode(input, return_tensors="pt")
num_tokens = input_tokens.shape[1]
if num_tokens > max_tokens:
return f"Error: Input exceeds the max token limit of {max_tokens}."
# Ensure min/max lengths are within bounds
min_summary_length = max(10, int(num_tokens * (min_length / 100)))
max_summary_length = min(max_tokens, int(num_tokens * (max_length / 100)))
# Summarize the input text
output = summarizer(input, min_length=min_summary_length, max_length=max_summary_length, truncation=True)
return output[0]['summary_text']
except Exception as e:
return f"Summarization failed: {str(e)}"
# Gradio Interface
with gr.Blocks() as demo:
with gr.Row():
model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="sshleifer/distilbart-cnn-12-6")
load_button = gr.Button("Load Model")
load_message = gr.Textbox(label="Load Status", interactive=False)
min_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Minimum Summary Length (%)", value=10)
max_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Maximum Summary Length (%)", value=20)
input_text = gr.Textbox(label="Input text to summarize", lines=6, value=example_text)
summarize_button = gr.Button("Summarize Text")
output_text = gr.Textbox(label="Summarized text", lines=4)
load_button.click(fn=load_model, inputs=model_dropdown, outputs=load_message)
summarize_button.click(fn=summarize_text, inputs=[input_text, min_length_slider, max_length_slider],
outputs=output_text)
demo.launch()
|