Spaces:
Runtime error
Runtime error
| 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() |