Walid Ahmed
commited on
Create app.py
Browse files
app.py
ADDED
|
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
import gradio as gr
|
| 3 |
+
from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer, AutoConfig
|
| 4 |
+
|
| 5 |
+
# List of summarization models
|
| 6 |
+
model_names = [
|
| 7 |
+
"google/bigbird-pegasus-large-arxiv",
|
| 8 |
+
"facebook/bart-large-cnn",
|
| 9 |
+
"google/t5-v1_1-large",
|
| 10 |
+
"sshleifer/distilbart-cnn-12-6",
|
| 11 |
+
"allenai/led-base-16384",
|
| 12 |
+
"google/pegasus-xsum",
|
| 13 |
+
"togethercomputer/LLaMA-2-7B-32K"
|
| 14 |
+
]
|
| 15 |
+
|
| 16 |
+
# Placeholder for the summarizer pipeline, tokenizer, and maximum tokens
|
| 17 |
+
summarizer = None
|
| 18 |
+
tokenizer = None
|
| 19 |
+
max_tokens = None
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
# Function to load the selected model
|
| 23 |
+
def load_model(model_name):
|
| 24 |
+
global summarizer, tokenizer, max_tokens
|
| 25 |
+
try:
|
| 26 |
+
# Load the summarization pipeline with the selected model
|
| 27 |
+
summarizer = pipeline("summarization", model=model_name, torch_dtype=torch.bfloat16)
|
| 28 |
+
# Load the tokenizer for the selected model
|
| 29 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
| 30 |
+
# Load the configuration for the selected model
|
| 31 |
+
config = AutoConfig.from_pretrained(model_name)
|
| 32 |
+
|
| 33 |
+
# Determine the maximum tokens based on available configuration attributes
|
| 34 |
+
if hasattr(config, 'max_position_embeddings'):
|
| 35 |
+
max_tokens = config.max_position_embeddings
|
| 36 |
+
elif hasattr(config, 'n_positions'):
|
| 37 |
+
max_tokens = config.n_positions
|
| 38 |
+
elif hasattr(config, 'd_model'):
|
| 39 |
+
max_tokens = config.d_model # for T5 models, d_model is a rough proxy
|
| 40 |
+
else:
|
| 41 |
+
max_tokens = "Unknown"
|
| 42 |
+
|
| 43 |
+
return f"Model {model_name} loaded successfully! Max tokens: {max_tokens}"
|
| 44 |
+
except Exception as e:
|
| 45 |
+
return f"Failed to load model {model_name}. Error: {str(e)}"
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
# Function to summarize the input text
|
| 49 |
+
def summarize_text(input, min_length, max_length):
|
| 50 |
+
if summarizer is None:
|
| 51 |
+
return "No model loaded!"
|
| 52 |
+
|
| 53 |
+
# Tokenize the input text and check the number of tokens
|
| 54 |
+
input_tokens = tokenizer.encode(input, return_tensors="pt")
|
| 55 |
+
num_tokens = input_tokens.shape[1]
|
| 56 |
+
if num_tokens > max_tokens:
|
| 57 |
+
# Return an error message if the input text exceeds the maximum token limit
|
| 58 |
+
return f"Error: The input text has {num_tokens} tokens, which exceeds the maximum allowed {max_tokens} tokens. Please enter shorter text."
|
| 59 |
+
|
| 60 |
+
# Calculate minimum and maximum summary length based on the percentages
|
| 61 |
+
min_summary_length = int(num_tokens * (min_length / 100))
|
| 62 |
+
max_summary_length = int(num_tokens * (max_length / 100))
|
| 63 |
+
|
| 64 |
+
# Summarize the input text using the loaded model with specified lengths
|
| 65 |
+
output = summarizer(input, min_length=min_summary_length, max_length=max_summary_length)
|
| 66 |
+
return output[0]['summary_text']
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
# Gradio Interface
|
| 70 |
+
with gr.Blocks() as demo:
|
| 71 |
+
with gr.Row():
|
| 72 |
+
# Dropdown menu for selecting the model
|
| 73 |
+
model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="sshleifer/distilbart-cnn-12-6")
|
| 74 |
+
# Button to load the selected model
|
| 75 |
+
load_button = gr.Button("Load Model")
|
| 76 |
+
|
| 77 |
+
# Textbox to display the load status
|
| 78 |
+
load_message = gr.Textbox(label="Load Status", interactive=False)
|
| 79 |
+
|
| 80 |
+
# Slider for minimum summary length
|
| 81 |
+
min_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Minimum Summary Length (%)", value=10)
|
| 82 |
+
# Slider for maximum summary length
|
| 83 |
+
max_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Maximum Summary Length (%)", value=20)
|
| 84 |
+
|
| 85 |
+
# Textbox for inputting the text to be summarized
|
| 86 |
+
input_text = gr.Textbox(label="Input text to summarize", lines=6)
|
| 87 |
+
# Button to trigger the summarization
|
| 88 |
+
summarize_button = gr.Button("Summarize Text")
|
| 89 |
+
# Textbox to display the summarized text
|
| 90 |
+
output_text = gr.Textbox(label="Summarized text", lines=4)
|
| 91 |
+
|
| 92 |
+
# Define the actions for the load button and summarize button
|
| 93 |
+
load_button.click(fn=load_model, inputs=model_dropdown, outputs=load_message)
|
| 94 |
+
summarize_button.click(fn=summarize_text, inputs=[input_text, min_length_slider, max_length_slider],
|
| 95 |
+
outputs=output_text)
|
| 96 |
+
|
| 97 |
+
# Launch the Gradio interface
|
| 98 |
+
demo.launch()
|