Update tasks/text.py

tasks/text.py

@@ -95,12 +95,12 @@ async def evaluate_text(request: TextEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
     #--------------------------------------------------------------------------------------------   
-    print('Start Binary')
+    logger.info('Start Binary')
     # Binary Model
-    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_BINARY)
+    tokenizer = AutoTokenizer.from_pretrained(BINARY_MODEL)
     print('Loaded Tokenizer')
     model = AutoModelForSequenceClassification.from_pretrained(BINARY_MODEL)
-    print('Loaded Model')
+    logger.info('Loaded Model')
     model.to(device)
     model.eval()
 
@@ -117,18 +117,18 @@ async def evaluate_text(request: TextEvaluationRequest):
             prediction = "0_not_relevant" if binary_prediction==0 else 1
             predictions.append(prediction)
         if i%10:
-            print(f'iteration: {i}')
+            logger.info(f'iteration: {i}')
             
     gc.collect()
 
     ## 2. Taxonomy Model
-    print('Start Multi')
-    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_MULTI_CLASS)
-    print('Loaded Tokenizer')
+    logger.info('Start Multi')
+    tokenizer = AutoTokenizer.from_pretrained(MULTI_CLASS_MODEL)
+    logger.info('Loaded Tokenizer')
     model = AutoModelForSequenceClassification.from_pretrained(MULTI_CLASS_MODEL)
     model.to(device)
     model.eval()
-    print('Loaded Model')
+    logger.info('Loaded Model')
     for i,text in tqdm(enumerate(test_dataset["quote"])):
         if isinstance(predictions[i], str):
             continue
@@ -142,7 +142,7 @@ async def evaluate_text(request: TextEvaluationRequest):
             prediction = ID2LABEL[taxonomy_prediction]
             predictions[i] = prediction
         if i%10:
-            print(f'iteration: {i}')
+            logger.info(f'iteration: {i}')
     predictions = [LABEL_MAPPING[pred] for pred in predictions]
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE