Experimenting with few-shot classification.

app.py

@@ -40,6 +40,17 @@ class OnnxSetFitModel:
         )
         return self.model_head.predict(embeddings.cpu())
 
+    def predict_proba(self, inputs):
+        encoded_inputs = self.tokenizer(
+            inputs, padding=True, truncation=True, return_tensors="pt"
+        ).to(self.ort_model.device)
+
+        outputs = self.ort_model(**encoded_inputs)
+        embeddings = mean_pooling(
+            outputs["last_hidden_state"], encoded_inputs["attention_mask"]
+        )
+        return self.model_head.predict_proba(embeddings.cpu())
+
     def __call__(self, inputs):
         return self.predict(inputs)
 
@@ -100,7 +111,9 @@ def zero_shot_classification(data):
 
 def few_shot_classification(data):
     results = onnx_few_shot_model(data['sequence'])
+    probs = onnx_few_shot_model.predict_proba(data['sequence'])
     print(results)
+    print(probs)
     response_string = json.dumps(results.tolist())
     return response_string