Spaces:
Runtime error
Runtime error
| import os | |
| from typing import Any, Callable, List, Optional, Tuple | |
| import nltk | |
| nltk.download('punkt') | |
| import gradio as gr | |
| from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM | |
| print(gr.__version__) | |
| # A folderpath for where the examples are stored | |
| EXAMPLES_FOLDER_NAME = "examples" | |
| # A List of repo names for the huggingface models available for inference | |
| HF_MODELS = ["huggingface/facebook/bart-large-cnn", | |
| "huggingface/sshleifer/distilbart-xsum-12-6", | |
| "huggingface/google/pegasus-xsum", | |
| "huggingface/philschmid/bart-large-cnn-samsum", | |
| "huggingface/linydub/bart-large-samsum", | |
| "huggingface/philschmid/distilbart-cnn-12-6-samsum", | |
| "huggingface/knkarthick/MEETING-SUMMARY-BART-LARGE-XSUM-SAMSUM-DIALOGSUM-AMI", | |
| ] | |
| ################################################################################ | |
| # Functions: Document statistics | |
| ################################################################################ | |
| # Function that uses a huggingface tokenizer to count how many tokens are in a text | |
| def count_tokens(input_text, model_path='sshleifer/distilbart-cnn-12-6'): | |
| # Load a huggingface tokenizer | |
| tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| # Tokenize the text | |
| tokens = tokenizer(input_text) | |
| # Count the number of tokens | |
| return len(tokens['input_ids']) | |
| # Function that uses nltk to count sentences in a text | |
| def count_sentences(input_text): | |
| # Use nltk to count sentences in the text | |
| number_of_sentences = nltk.sent_tokenize(input_text) | |
| # Return the number of sentences | |
| return len(number_of_sentences) | |
| # Function that counts the number of words in a text | |
| def count_words(input_text): | |
| # Use nltk to count words in the text | |
| number_of_words = nltk.word_tokenize(input_text) | |
| # Return the number of words | |
| return len(number_of_words) | |
| # Function that computes a few document statistics such as the number of tokens, sentences, and words | |
| def compute_stats(input_text, models: Optional[List[str]] = None): | |
| # Count the number of tokens | |
| num_tokens = count_tokens(input_text) | |
| # Count the number of sentences | |
| num_sentences = count_sentences(input_text) | |
| # Count the number of words | |
| num_words = count_words(input_text) | |
| # Return the document statistics formatted as a string | |
| output_str = "| Tokens: {0} \n| Sentences: {1} \n| Words: {2}".format(num_tokens, num_sentences, num_words) + "\n" | |
| output_str += "The max number of tokens for the model is: 1024" + "\n" # I manually set 1024 as the max. I don't intend to use any models that are smaller anyway. | |
| # output_str += "Number of documents splits: 17.5" | |
| return output_str | |
| # # A function to loop through a list of strings | |
| # # returning the last element in the filepath for each string | |
| # def get_file_names(file_paths): | |
| # # Create a list of file names | |
| # file_names = [] | |
| # # Loop through the file paths | |
| # for file_path in file_paths: | |
| # # Get the last element in the file path | |
| # file_name = file_path.split('/')[-2:] | |
| # # Add the file name to the list | |
| # file_names.append(file_name) | |
| # # Loop through the file names and append to a string | |
| # file_names_str = "" | |
| # for file_name in file_names: | |
| # breakpoint() | |
| # file_names_str += file_name[0] + "\n" | |
| # # Return the list of file names | |
| # return file_names_str | |
| ################################################################################ | |
| # Functions: Huggingface Inference | |
| ################################################################################ | |
| # Function that uses a huggingface pipeline to predict a summary of a text | |
| # input is a text string of a dialog conversation | |
| def predict(dialog_text): | |
| # Load a huggingface model | |
| model = pipeline('summarization', model="philschmid/bart-large-cnn-samsum") #model='sshleifer/distilbart-cnn-12-6') | |
| # Build tokenizer_kwargs to set a max length and truncate the data on inference | |
| tokenizer_kwargs = {'truncation': True, 'max_length': 1024} | |
| # Use the model to predict a summary of the text | |
| summary = model(dialog_text, **tokenizer_kwargs) | |
| # Return the summary w/ the model name | |
| output = f"{hf_model_name} output: {summary[0]['summary_text']}" | |
| return output, "output2" | |
| def recursive_predict(dialog_text: str, hf_model_name: Tuple[str]): | |
| breakpoint() | |
| asdf = "asdf" | |
| return output | |
| ################################################################################ | |
| # Functions: Gradio Utilities | |
| ################################################################################ | |
| # Function to build examples for gradio app | |
| # Load text files from the examples folder as a list of strings for gradio | |
| def get_examples(folder_path): | |
| # Create a list of strings | |
| examples = [] | |
| # Loop through the files in the folder | |
| for file in os.listdir(folder_path): | |
| # Load the file | |
| with open(os.path.join(folder_path, file), 'r') as f: | |
| # Add the file to the list | |
| examples.append([f.read(), ["None"]]) | |
| # Return the list of strings | |
| return examples | |
| # A function that loops through a list of model paths, creates a gradio interface with the | |
| # model name, and adds it to the list of interfaces | |
| # It outputs a list of interfaces | |
| def get_hf_interfaces(models_to_load): | |
| # Create a list of interfaces | |
| interfaces = [] | |
| # Loop through the HF_MODELS | |
| for model in models_to_load: | |
| # Create a gradio interface with the model name | |
| interface = gr.Interface.load(model, title="this is a test TITLE", alias="this is an ALIAS") | |
| # Add the interface to the list | |
| interfaces.append(interface) | |
| # Return the list of interfaces | |
| return interfaces | |
| ################################################################################ | |
| # Build Gradio app | |
| ################################################################################ | |
| # print_details = gr.Interface( | |
| # fn=lambda x: get_file_names(HF_MODELS), | |
| # inputs="text", | |
| # outputs="text", | |
| # title="Statistics of the document" | |
| # ) | |
| # Outputs a string of various document statistics | |
| document_statistics = gr.Interface( | |
| fn=compute_stats, | |
| inputs="text", | |
| outputs="text", | |
| title="Statistics of the document" | |
| ) | |
| maddie_mixer_summarization = gr.Interface( | |
| fn=recursive_predict, | |
| inputs="text", | |
| outputs="text", | |
| title="Statistics of the document" | |
| ) | |
| # Build Examples to pass along to the gradio app | |
| examples = get_examples(EXAMPLES_FOLDER_NAME) | |
| # Build a list of huggingface interfaces from model paths, | |
| # then add document statistics, and any custom interfaces | |
| all_interfaces = get_hf_interfaces(HF_MODELS) | |
| all_interfaces.insert(0, document_statistics) # Insert the statistics interface at the beginning | |
| # all_interfaces.insert(0, print_details) | |
| # all_interfaces.append(maddie_mixer_summarization) # Add the interface for the maddie mixer | |
| # Build app | |
| app = gr.Parallel(*all_interfaces, | |
| title='Text Summarizer (Maddie Custom)', | |
| description="Write a summary of a text", | |
| # examples=examples, | |
| inputs=gr.inputs.Textbox(lines = 10, label="Text"), | |
| ) | |
| # Launch | |
| app.launch(inbrowser=True, show_error=True) |