Introduced support for sentiment analysis (#81)

* Create nlp.py & add_sentiment(df) function

* Added optional NLP dependencies

* Closes #60 

---------

Co-authored-by: joweich <[email protected]>

.gitignore

@@ -151,3 +151,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# VScode
+.vscode/

chatminer/nlp.py

@@ -0,0 +1,46 @@
+from transformers import pipeline
+import pandas as pd
+
+
+def add_sentiment(df: pd.DataFrame, lang: str = "en") -> pd.DataFrame:
+    """
+    Add sentiment column to the input dataframe
+
+    Parameters:
+    df (pd.DataFrame): The input dataframe
+    lang (str): Language of the messages, defaults to "en"
+
+    Returns:
+    pd.DataFrame: The input dataframe with an additional column "sentiment"
+
+    """
+    if "message" not in df.columns:
+        raise ValueError("Input dataframe does not contain a 'message' column")
+
+    model_path = (
+        "cardiffnlp/twitter-roberta-base-sentiment-latest"
+        if lang == "en"
+        else "cardiffnlp/twitter-xlm-roberta-base-sentiment"
+    )
+    sentiment_pipeline = pipeline("sentiment-analysis", model=model_path)
+
+    def extract_sentiment(message: str) -> str:
+        """
+        Extract sentiment from message
+
+        Parameters:
+        message (str): The input message
+
+        Returns:
+        str: The sentiment of the message
+
+        """
+        try:
+            return sentiment_pipeline(message)[0]["label"]
+        except:
+            print(f"Error processing message: {message}")
+            return None
+
+    df["sentiment"] = df["message"].apply(extract_sentiment)
+
+    return df

requirements.txt

@@ -3,4 +3,5 @@ numpy
 matplotlib
 wordcloud
 python-dateutil
-tqdm
+tqdm
+transformers

setup.cfg

@@ -38,3 +38,7 @@ install_requires =
     wordcloud
     python-dateutil
     tqdm
+
+[options.extras_require]
+NLP =
+    transformers