Update app.py

app.py

@@ -1 +1,83 @@
-print("hi")
+from ast import Str
+import gradio as gr
+from tweetnlp import Sentiment, NER
+from typing import Tuple, Dict
+from statistics import mean
+
+def clean_tweet(tweet: str, remove_chars: str = "@#") -> str:
+    """Remove any unwanted characters
+    Args:
+        tweet (str): The raw tweet
+        remove_chars (str, optional): The characters to remove. Defaults to "@#".
+    Returns:
+        str: The tweet with these characters removed
+    """
+    for char in remove_chars:
+        tweet = tweet.replace(char, "")
+    return tweet
+
+
+def format_sentiment(model_output: Dict) -> Dict:
+    """Format the output of the sentiment model
+    Args:
+        model_output (Dict): The model output
+    Returns:
+        Dict: The format for gradio
+    """
+    formatted_output = dict()
+    if model_output["label"] == "positive":
+        formatted_output["positive"] = model_output["probability"]
+        formatted_output["negative"] = 1 - model_output["probability"]
+    else:
+        formatted_output["negative"] = model_output["probability"]
+        formatted_output["positive"] = 1 - model_output["probability"]
+    return formatted_output
+
+
+def format_entities(model_output: Dict) -> Dict:
+    """Format the output of the NER model
+    Args:
+        model_output (Dict): The model output
+    Returns:
+        Dict: The format for gradio
+    """
+    formatted_output = dict()
+    for entity in model_output["entity_prediction"]:
+        new_output = dict()
+        name = " ".join(entity["entity"])
+        entity_type = entity["type"]
+        new_key = f"{name}:{entity_type}"
+        new_value = mean(entity["probability"])
+        formatted_output[new_key] = new_value
+    return formatted_output
+
+
+def classify(tweet: str) -> Tuple[Dict, Dict]:
+    """Runs models
+    Args:
+        tweet (str): The raw tweet
+    Returns:
+        Tuple[Dict, Dict]: The formatted_sentiment and formatted_entities of the tweet
+    """
+    tweet = clean_tweet(tweet)
+    # Get sentiment
+    model_sentiment = se_model.sentiment(tweet)
+    formatted_sentiment = format_sentiment(model_sentiment)
+    # Get entities
+    entities = ner_model.ner(tweet)
+    formatted_entities = format_entities(entities)
+    return formatted_sentiment, formatted_entities
+
+
+if __name__ == "__main__":
+    # https://github.com/cardiffnlp/tweetnlp
+    se_model = Sentiment()
+    ner_model = NER()
+
+    # Get a few examples from: https://twitter.com/NFLFantasy
+    examples = list()
+    examples.append("Dameon Pierce is clearly the #Texans starter and he once again looks good")
+    examples.append("Deebo Samuel had 150+ receiving yards in 4 games last year - the most by any receiver in the league.")
+
+    iface = gr.Interface(fn=classify, inputs="text", outputs=["label", "label"], examples=examples)
+    iface.launch()