Walid Ahmed commited on
Commit
cee775f
1 Parent(s): e1a3f5e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +98 -0
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()