Spaces:
Running
Running
| import fasttext | |
| import gradio as gr | |
| import joblib | |
| import json as js | |
| import omikuji | |
| import os | |
| import re | |
| from collections import defaultdict | |
| from huggingface_hub import snapshot_download | |
| from typing import List, Tuple, Dict | |
| from install_packages import download_model | |
| download_model('https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin', 'lid.176.bin') | |
| # Download the model files from Hugging Face | |
| for repo_id in ['kapllan/omikuji-bonsai-parliament-de-spacy', 'kapllan/omikuji-bonsai-parliament-fr-spacy', | |
| 'kapllan/omikuji-bonsai-parliament-it-spacy']: | |
| if not os.path.exists(repo_id): | |
| os.makedirs(repo_id) | |
| model_dir = snapshot_download(repo_id=repo_id, local_dir=repo_id) | |
| lang_model = fasttext.load_model('lid.176.bin') | |
| with open('./id2label.json', 'r') as f: | |
| id2label = js.load(f) | |
| with open('topics_hierarchy.json', 'r') as f: | |
| topics_hierarchy = js.load(f) | |
| def map_language(language: str) -> str: | |
| language_mapping = {'de': 'German', | |
| 'it': 'Italian', | |
| 'fr': 'French'} | |
| if language in language_mapping.keys(): | |
| return language_mapping[language] | |
| else: | |
| return language | |
| def find_model(language: str): | |
| vectorizer, model = None, None | |
| if language in ['de', 'fr', 'it']: | |
| path_to_vectorizer = f'./kapllan/omikuji-bonsai-parliament-{language}-spacy/vectorizer' | |
| path_to_model = f'./kapllan/omikuji-bonsai-parliament-{language}-spacy/omikuji-model' | |
| vectorizer = joblib.load(path_to_vectorizer) | |
| model = omikuji.Model.load(path_to_model) | |
| return vectorizer, model | |
| def predict_lang(text: str) -> str: | |
| text = re.sub(r'\n', '', text) # Remove linebreaks because fasttext cannot process that otherwise | |
| predictions = lang_model.predict(text, k=1) # returns top 2 matching languages | |
| language = predictions[0][0] # returns top 2 matching languages | |
| language = re.sub(r'__label__', '', language) # returns top 2 matching languages | |
| return language | |
| def predict_topic(text: str) -> [List[str], str]: | |
| results = [] | |
| language = predict_lang(text) | |
| vectorizer, model = find_model(language) | |
| language = map_language(language) | |
| if vectorizer is not None: | |
| texts = [text] | |
| vector = vectorizer.transform(texts) | |
| for row in vector: | |
| if row.nnz == 0: # All zero vector, empty result | |
| continue | |
| feature_values = [(col, row[0, col]) for col in row.nonzero()[1]] | |
| for subj_id, score in model.predict(feature_values, top_k=1000): | |
| results.append((id2label[str(subj_id)], score)) | |
| return results, language | |
| def get_row_color(type: str): | |
| if 'main' in type.lower(): | |
| return 'background-color: darkgrey;' | |
| if 'sub' in type.lower(): | |
| return 'background-color: lightgrey;' | |
| def generate_html_table(topics: List[Tuple[str, str, float]]): | |
| html = '<table style="width:100%; border: 1px solid black; border-collapse: collapse;">' | |
| html += '<tr><th>Type</th><th>Topic</th><th>Score</th></tr>' | |
| for type, topic, score in topics: | |
| color = get_row_color(type) | |
| topic = f"<strong>{topic}</strong>" if 'main' in type.lower() else topic | |
| type = f"<strong>{type}</strong>" if 'main' in type.lower() else type | |
| score = f"<strong>{score}</strong>" if 'main' in type.lower() else score | |
| html += f'<tr style="{color}"><td>{type}</td><td>{topic}</td><td>{score}</td></tr>' | |
| html += '</table>' | |
| return html | |
| def restructure_topics(topics: List[Tuple[str, float]]) -> List[Tuple[str, str, float]]: | |
| topics = [(str(x[0]).lower(), x[1]) for x in topics] | |
| topics_as_dict = {} | |
| for predicted_topic, score in topics: | |
| if str(predicted_topic).lower() in topics_hierarchy.keys(): | |
| topics_as_dict[str(predicted_topic).lower()] = [] | |
| for predicted_topic, score in topics: | |
| for main_topic, sub_topics in topics_hierarchy.items(): | |
| if main_topic in topics_as_dict.keys() and predicted_topic != main_topic and predicted_topic in sub_topics: | |
| topics_as_dict[main_topic].append(predicted_topic) | |
| topics_restructured = [] | |
| for predicted_main_topic, predicted_sub_topics in topics_as_dict.items(): | |
| if len(predicted_sub_topics) > 0: | |
| score = [t for t in topics if t[0] == predicted_main_topic][0][1] | |
| topics_restructured.append( | |
| ('Main Topic', predicted_main_topic, score)) | |
| predicted_sub_topics_with_scores = [] | |
| for pst in predicted_sub_topics: | |
| score = [t for t in topics if t[0] == pst][0][1] | |
| entry = ('Sub Topic', pst, score) | |
| if entry not in predicted_sub_topics_with_scores: | |
| predicted_sub_topics_with_scores.append(entry) | |
| for x in predicted_sub_topics_with_scores: | |
| topics_restructured.append(x) | |
| return topics_restructured | |
| def topic_modeling(text: str, threshold: float) -> [List[str], str]: | |
| # Prepare labels and scores for the plot | |
| sorted_topics, language = predict_topic(text) | |
| if len(sorted_topics) > 0 and language in ['German', 'French', 'Italian']: | |
| sorted_topics = [t for t in sorted_topics if t[1] >= threshold] | |
| else: | |
| sorted_topics = [] | |
| sorted_topics = restructure_topics(sorted_topics) | |
| sorted_topics = generate_html_table(sorted_topics) | |
| return sorted_topics, language | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# Topic Modeling") | |
| gr.Markdown("Enter a document and get each topic along with its score.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_text = gr.Textbox(lines=10, placeholder="Enter a document") | |
| submit_button = gr.Button("Submit") | |
| threshold_slider = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label="Score Threshold", | |
| value=0.0) | |
| language_text = gr.Textbox(lines=1, placeholder="Detected language will be shown here...", | |
| interactive=False, label="Detected Language") | |
| with gr.Column(): | |
| output_data = gr.HTML() | |
| submit_button.click(topic_modeling, inputs=[input_text, threshold_slider], | |
| outputs=[output_data, language_text]) | |
| # Launch the app | |
| iface.launch(share=True) | |