Spaces:
Runtime error
Runtime error
add show emotions name
Browse files
app.py
CHANGED
@@ -14,7 +14,8 @@ def translate(text):
|
|
14 |
logits = output.logits
|
15 |
prediction = F.softmax(logits, dim=1)
|
16 |
y_pred = torch.argmax(prediction).numpy()
|
17 |
-
|
|
|
18 |
# output = model.generate(input["input_ids"], max_length=40, num_beams=4, early_stopping=True)
|
19 |
|
20 |
title = "Text Emotion Classification"
|
|
|
14 |
logits = output.logits
|
15 |
prediction = F.softmax(logits, dim=1)
|
16 |
y_pred = torch.argmax(prediction).numpy()
|
17 |
+
class_names = ['sad','joy','love','anger','fear','surprise']
|
18 |
+
return class_names[y_pred]
|
19 |
# output = model.generate(input["input_ids"], max_length=40, num_beams=4, early_stopping=True)
|
20 |
|
21 |
title = "Text Emotion Classification"
|