update

tasks/text.py

@@ -1,7 +1,7 @@
 from fastapi import APIRouter
 from datetime import datetime
 from datasets import load_dataset, Dataset
-from sklearn.metrics import accuracy_score
+from sklearn.metrics import accuracy_score, f1_score
 import random
 from torch.utils.data import DataLoader
 
@@ -64,6 +64,13 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------   
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model_name = ["cococli/bert-base-uncased-frugalai", 'cococli/roberta-base-frugalai', "cococli/distilbert-base-uncased-frugalai",
+                    "cococli/albert-base-v2-frugalai", "cococli/bert-base-uncased-coco-frugalai",
+                     "cococli/distilbert-base-uncased-coco-frugalai", "cococli/albert-base-v2-coco-frugalai","cococli/electra-small-discriminator-coco-frugalai", 
+                     'cococli/roberta-base-coco-frugalai', "cococli/distilbert-base-uncased-climate-frugalai","cococli/albert-base-v2-climate-frugalai", 
+                     "cococli/electra-small-discriminator-frugalai", "cococli/bert-base-uncased-climate-frugalai","cococli/roberta-base-climate-frugalai",
+    ]
+
     
     tokenizer = AutoTokenizer.from_pretrained("cococli/bert-base-uncased-frugalai")
     model = AutoModelForSequenceClassification.from_pretrained("cococli/bert-base-uncased-frugalai").to(device)
@@ -100,7 +107,7 @@ async def evaluate_text(request: TextEvaluationRequest):
             # Forward pass through the model
             p = model(**tokenized_inputs)
             output = torch.argmax(p.logits, dim=1).cpu().numpy()
-            print(p)
+            # print(p)
             predictions = np.append(predictions, output)
 
     print("Finished prediction run")
@@ -119,6 +126,7 @@ async def evaluate_text(request: TextEvaluationRequest):
     print('Accuracy: ', (true_labels == predictions)/len(true_labels))
 
     print('Accuracy: ', accuracy_score(true_labels, predictions))
+    print('F1 SCORE: ', f1_score(true_labels, predictions))
     # Stop tracking emissions
     emissions_data = tracker.stop_task()