Create app.py

app.py

@@ -0,0 +1,144 @@
+import gradio as gr
+import gc
+from collections import defaultdict
+import torch
+from transformers import pipeline
+from lingua import Language, LanguageDetectorBuilder
+
+default_models = {
+    Language.ENGLISH: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.JAPANESE: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.ARABIC: "Ammar-alhaj-ali/arabic-MARBERT-sentiment",
+    Language.GERMAN: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.SPANISH: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.FRENCH: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.CHINESE: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.INDONESIAN: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.HINDI: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.ITALIAN: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.MALAY: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.PORTUGUESE: "lxyuan/distilbert-base-multilingual-cased-sentiments-student",
+    Language.SWEDISH: "KBLab/robust-swedish-sentiment-multiclass",
+    Language.FINNISH: "fergusq/finbert-finnsentiment",
+}
+language_detector = LanguageDetectorBuilder.from_all_languages().build()
+
+
+def split_message(message, max_length):
+    """ Split a message into a list of chunks of given maximum size. """
+    return [message[i: i + max_length] for i in range(0, len(message), max_length)]
+
+
+def process_messages_in_batches(messages_with_languages, models=None, max_length=512):
+    """
+    Process messages in batches, creating only one pipeline at a time, and maintain the original order.
+    Params:
+    messages_with_languages: list of tuples, each containing a message and its detected language
+    models: dict, model paths indexed by Language
+    Returns:
+    OrderedDict: containing the index as keys and tuple of (message, sentiment result) as values
+    """
+
+    if models is None:
+        models = default_models
+    else:
+        models = default_models.copy().update(models)
+
+    results = {}
+
+    # Group messages by model, preserving original order.
+    # If language is no detected or a model for that language is not
+    # provided, add None to results
+    messages_by_model = defaultdict(list)
+    for index, (message, language) in enumerate(messages_with_languages):
+        model_name = models.get(language)
+        if model_name:
+            messages_by_model[model_name].append((index, message))
+        else:
+            results[index] = {"label": "none", "score": 0}
+
+    # Process messages and maintain original order
+    for model_name, batch in messages_by_model.items():
+        sentiment_pipeline = pipeline(model=model_name)
+
+        chunks = []
+        message_map = {}
+        for idx, message in batch:
+            message_chunks = split_message(message, max_length)
+            for chunk in message_chunks:
+                chunks.append(chunk)
+                if idx in message_map:
+                    message_map[idx].append(len(chunks) - 1)
+                else:
+                    message_map[idx] = [len(chunks) - 1]
+
+        chunk_sentiments = sentiment_pipeline(chunks)
+
+        for idx, chunk_indices in message_map.items():
+            sum_scores = {"neutral": 0}
+            for chunk_idx in chunk_indices:
+                label = chunk_sentiments[chunk_idx]["label"]
+                score = chunk_sentiments[chunk_idx]["score"]
+                if label in sum_scores:
+                    sum_scores[label] += score
+                else:
+                    sum_scores[label] = score
+            best_sentiment = max(sum_scores, key=sum_scores.get)
+            score = sum_scores[best_sentiment] / len(chunk_indices)
+            results[idx] = {"label": best_sentiment, "score": score}
+
+        # Force garbage collections to remove the model from memory
+        del sentiment_pipeline
+        gc.collect()
+
+    # Unify common spellings of the labels
+    for i in range(len(results)):
+        results[i]["label"] = results[i]["label"].lower()
+
+    results = [results[i] for i in range(len(results))]
+
+    return results
+
+
+def sentiment(messages, models=None):
+    """
+    Estimate the sentiment of a list of messages (strings of text). The
+    sentences may be in different languages from each other.
+    We maintain a list of default models for some languages. In addition,
+    the user can provide a model for a given language in the models
+    dictionary. The keys for this dictionary are lingua.Language objects
+    and items HuggingFace model paths.
+    Params:
+    messages: list of message strings
+    models: dict, huggingface model paths indexed by lingua.Language
+    Returns:
+    OrderedDict: containing the index as keys and tuple of (message, sentiment result) as values
+    """
+    messages_with_languages = [
+        (message, language_detector.detect_language_of(message)) for message in messages
+    ]
+
+    results = process_messages_in_batches(messages_with_languages, models)
+    return results
+
+
+def main(input_text):
+    messages = [message.strip() for message in input_text.split('\n') if message.strip()]
+    results = sentiment(messages)
+    output = []
+    for idx, result in enumerate(results):
+        message = messages[idx]
+        sentiment_label = result["label"]
+        sentiment_score = result["score"]
+        output.append((message, sentiment_label, sentiment_score))
+    return output
+
+
+iface = gr.Interface(
+    fn=main,
+    inputs="text",
+    outputs=[gr.outputs.Table(headings=["Message", "Sentiment", "Score"])],
+    title="Sentiment Analysis Pipeline",
+    description="Enter your messages (one per line) and get sentiment analysis results.",
+)
+iface.launch()