Add emotions model

.gitignore

@@ -1,3 +1,5 @@
+example_notebook.ipynb
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

main.py

@@ -1,6 +1,9 @@
-from fastapi import FastAPI
-from yt_api import get_comments
+from fastapi import FastAPI, Response
 from pydantic_settings import BaseSettings, SettingsConfigDict
+import pandas as pd
+
+from yt_api import get_comments
+from models import init_emotions_model
 
 
 class Settings(BaseSettings):
@@ -11,8 +14,7 @@ class Settings(BaseSettings):
 settings = Settings()
 app = FastAPI(title='social-stat')
 
-
-YT_API_KEY = settings.YT_API_KEY
+emotions_clf = init_emotions_model()
 
 
 @app.get('/')
@@ -20,6 +22,27 @@ def home():
     return 'social-stat'
 
 
-@app.post('/predict')
+@app.get('/predict')
 def predict(video_id):
-    return get_comments(video_id, YT_API_KEY)
+    # Get comments
+    comments = get_comments(video_id, settings.YT_API_KEY)
+    comments_df = pd.DataFrame(comments)
+
+    # Predict emotions
+    text_list = comments_df['text_display'].to_list()
+    preds = emotions_clf(text_list)
+
+    # Add predictions to DataFrame
+    preds_df = []
+    for pred in preds:
+        pred_dict = {}
+        for emotion in pred:
+            pred_dict[emotion['label']] = emotion['score']
+        preds_df.append(pred_dict)
+    preds_df = pd.DataFrame(preds_df)
+    comments_df = pd.concat([comments_df, preds_df], axis=1)
+
+    # Return DataFrame as a JSON file
+    return Response(
+        content=comments_df.to_json(orient='records'),
+        media_type='application/json')

models.py

@@ -0,0 +1,10 @@
+from transformers import pipeline
+
+
+def init_emotions_model():
+    classifier = pipeline(
+        task="text-classification",
+        model="SamLowe/roberta-base-go_emotions",
+        top_k=None)
+
+    return classifier

requirements.txt

@@ -2,3 +2,8 @@ requests
 fastapi
 uvicorn
 pydantic_settings
+torch
+torchvision
+torchaudio
+transformers
+pandas

yt_api.py

@@ -13,7 +13,7 @@ def get_comments(video_id, api_key):
     while 'nextPageToken' in response.keys():
         response = get_response(
             video_id, api_key, page_token=response['nextPageToken'])
-        comment_list.update(response_to_comments(response))
+        comment_list += (response_to_comments(response))
 
     return comment_list
 
@@ -34,7 +34,7 @@ def get_response(video_id, api_key, page_token=None, max_results=100):
 
 def response_to_comments(response):
     """Converts JSON response to `comment_list` dict."""
-    comment_list = {}
+    comment_list = []
     for comment in response['items']:
         comment = comment['snippet']
         can_reply = comment['canReply']
@@ -43,7 +43,8 @@ def response_to_comments(response):
         comment_id = comment['id']
         comment = comment['snippet']
         try:
-            comment_list[comment_id] = {
+            comment_list.append({
+                'comment_id': comment_id,
                 'video_id': comment['videoId'],
                 'channel_id': comment['authorChannelId']['value'],
                 'author_display_name': comment['authorDisplayName'],
@@ -54,7 +55,7 @@ def response_to_comments(response):
                 'like_count': comment['likeCount'],
                 'can_reply': can_reply,
                 'total_reply_count': total_reply_count,
-            }
+            })
         except Exception as e:
             print(f"Error: {e}\nComment: {comment}")
             continue