ceva

app.py

@@ -37,12 +37,13 @@ def ssl_predict(in_text, model_type):
         model = FixMatchTune(encoder_name="andrei-saceleanu/ro-offense-freematch")
         model.cls_head.load_weights("./checkpoints/freematch_tune")
 
-    print(toks["input_ids"])
     preds, _ = model([toks["input_ids"],toks["attention_mask"]], training=False)
-    print(preds)
     probs = list(preds[0].numpy())
 
-    return {k:v for k, v in zip(label_names, probs)}
+    d = {}
+    for k, v in zip(label_names, probs):
+        d[k] = float(v)
+    return d