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)