Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
# Example text for summarization
|
22 |
+
example_text = (
|
23 |
+
"Artificial intelligence (AI) is intelligence—perceiving, synthesizing, and inferring information—"
|
24 |
+
"demonstrated by machines, as opposed to intelligence displayed by non-human animals and humans. "
|
25 |
+
"Example tasks in which AI is employed include speech recognition, computer vision, language translation, "
|
26 |
+
"autonomous vehicles, and game playing. AI research has been defined as the field of study of intelligent "
|
27 |
+
"agents, which refers to any system that perceives its environment and takes actions that maximize its "
|
28 |
+
"chance of achieving its goals."
|
29 |
+
)
|
30 |
+
|
31 |
+
# Function to load the selected model
|
32 |
+
def load_model(model_name):
|
33 |
+
global summarizer, tokenizer, max_tokens
|
34 |
+
try:
|
35 |
+
# Load the summarization pipeline with the selected model
|
36 |
+
summarizer = pipeline("summarization", model=model_name, torch_dtype=torch.float32)
|
37 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
38 |
+
config = AutoConfig.from_pretrained(model_name)
|
39 |
+
|
40 |
+
# Set a reasonable default for max_tokens if not available
|
41 |
+
max_tokens = getattr(config, 'max_position_embeddings', 1024)
|
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 |
+
# Function to summarize the input text
|
48 |
+
def summarize_text(input, min_length, max_length):
|
49 |
+
if summarizer is None:
|
50 |
+
return "No model loaded!"
|
51 |
+
|
52 |
+
try:
|
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 f"Error: Input exceeds the max token limit of {max_tokens}."
|
58 |
+
|
59 |
+
# Ensure min/max lengths are within bounds
|
60 |
+
min_summary_length = max(10, int(num_tokens * (min_length / 100)))
|
61 |
+
max_summary_length = min(max_tokens, int(num_tokens * (max_length / 100)))
|
62 |
+
|
63 |
+
# Summarize the input text
|
64 |
+
output = summarizer(input, min_length=min_summary_length, max_length=max_summary_length, truncation=True)
|
65 |
+
return output[0]['summary_text']
|
66 |
+
except Exception as e:
|
67 |
+
return f"Summarization failed: {str(e)}"
|
68 |
+
|
69 |
+
# Gradio Interface
|
70 |
+
with gr.Blocks() as demo:
|
71 |
+
with gr.Row():
|
72 |
+
model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="sshleifer/distilbart-cnn-12-6")
|
73 |
+
load_button = gr.Button("Load Model")
|
74 |
+
|
75 |
+
load_message = gr.Textbox(label="Load Status", interactive=False)
|
76 |
+
|
77 |
+
min_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Minimum Summary Length (%)", value=10)
|
78 |
+
max_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Maximum Summary Length (%)", value=20)
|
79 |
+
|
80 |
+
input_text = gr.Textbox(label="Input text to summarize", lines=6, value=example_text)
|
81 |
+
summarize_button = gr.Button("Summarize Text")
|
82 |
+
output_text = gr.Textbox(label="Summarized text", lines=4)
|
83 |
+
|
84 |
+
load_button.click(fn=load_model, inputs=model_dropdown, outputs=load_message)
|
85 |
+
summarize_button.click(fn=summarize_text, inputs=[input_text, min_length_slider, max_length_slider],
|
86 |
+
outputs=output_text)
|
87 |
+
|
88 |
+
demo.launch()
|