Nice

tulane-cmps6730 · Apr 30, 2024 · e366dba · e366dba
1 parent adb1f35
commit e366dba
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nlp/app/routes.py b/nlp/app/routes.py
@@ -64,10 +64,10 @@ def index():
             probas = (preds > 0.5).astype(int) 
             if preds == 1:
                 prediction = "WIN"
-                proba = probas
+                proba = probas[1]
             else:
                 prediction = "LOSS"
-                proba = proba
+                proba = proba[0]
         elif model_choice == 'bert':
             tokenizer = BertTokenizerFast.from_pretrained('prajjwal1/bert-mini')
             model = AutoModelForSequenceClassification.from_pretrained('/Users/jackiecollopy/Downloads/project-reddit/notebooks/bert.pth')