Spaces:
Runtime error
Runtime error
| 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 | |
| # 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.bfloat16) | |
| # Load the tokenizer for the selected model | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| # Load the configuration for the selected model | |
| config = AutoConfig.from_pretrained(model_name) | |
| # Determine the maximum tokens based on available configuration attributes | |
| if hasattr(config, 'max_position_embeddings'): | |
| max_tokens = config.max_position_embeddings | |
| elif hasattr(config, 'n_positions'): | |
| max_tokens = config.n_positions | |
| elif hasattr(config, 'd_model'): | |
| max_tokens = config.d_model # for T5 models, d_model is a rough proxy | |
| else: | |
| max_tokens = "Unknown" | |
| 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!" | |
| # 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 an error message if the input text exceeds the maximum token limit | |
| return f"Error: The input text has {num_tokens} tokens, which exceeds the maximum allowed {max_tokens} tokens. Please enter shorter text." | |
| # Calculate minimum and maximum summary length based on the percentages | |
| min_summary_length = int(num_tokens * (min_length / 100)) | |
| max_summary_length = int(num_tokens * (max_length / 100)) | |
| # Summarize the input text using the loaded model with specified lengths | |
| output = summarizer(input, min_length=min_summary_length, max_length=max_summary_length) | |
| return output[0]['summary_text'] | |
| # Gradio Interface | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| # Dropdown menu for selecting the model | |
| model_dropdown = gr.Dropdown(choices=model_names, label="Choose a model", value="sshleifer/distilbart-cnn-12-6") | |
| # Button to load the selected model | |
| load_button = gr.Button("Load Model") | |
| # Textbox to display the load status | |
| load_message = gr.Textbox(label="Load Status", interactive=False) | |
| # Slider for minimum summary length | |
| min_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Minimum Summary Length (%)", value=10) | |
| # Slider for maximum summary length | |
| max_length_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Maximum Summary Length (%)", value=20) | |
| # Textbox for inputting the text to be summarized | |
| input_text = gr.Textbox(label="Input text to summarize", lines=6) | |
| # Button to trigger the summarization | |
| summarize_button = gr.Button("Summarize Text") | |
| # Textbox to display the summarized text | |
| output_text = gr.Textbox(label="Summarized text", lines=4) | |
| # Define the actions for the load button and summarize button | |
| 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) | |
| # Launch the Gradio interface | |
| demo.launch() | |