refactor: move NLP processing out of server file

This commit is contained in:
2026-02-05 17:10:48 +00:00
parent 8e257a92d9
commit d4db7bec24
2 changed files with 31 additions and 26 deletions

28
server/nlp_processor.py Normal file
View File

@@ -0,0 +1,28 @@
import torch
import pandas as pd
from transformers import pipeline
emotion_classifier = pipeline(
"text-classification",
model="j-hartmann/emotion-english-distilroberta-base",
top_k=None,
truncation=True,
device=0 if torch.cuda.is_available() else -1
)
def add_emotion_cols(df: pd.Dataframe, content_col: str) -> None:
texts = df[content_col].astype(str).str.slice(0, 512).tolist()
results = emotion_classifier(
texts,
batch_size=64
)
labels = [r["label"] for r in results[0]]
for label in labels:
df[f"emotion_{label}"] = [
next(item["score"] for item in row if item["label"] == label)
for row in results
]