Experimenting with few-shot classification.

app.py

@@ -71,7 +71,7 @@ ort_model = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5
 few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
 
 # Train few_shot_model
-candidate_labels = ["supported", "refuted"]
+candidate_labels = ["true", "false"]
 reference_dataset = load_dataset("emotion")
 dummy_dataset = Dataset.from_dict({})
 train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="Based on the Given passage, the hypothesis is {}.")
@@ -110,16 +110,14 @@ def zero_shot_classification(data):
     return response_string
 
 def create_sequences(data):
-    return ['###Given:\n' + data['sequence'] + '\n###End Given\n###Hypothesis:\n' + data['hypothesis_template'].format(label) + "\n###End Hypothesis" + "\nBased on the Given passage, the Hypothesis is {}" for label in data['candidate_labels']]
+    return ['###Given:\n' + data['sequence'] + '\n###End Given\n###Hypothesis:\n' + data['hypothesis_template'].format(label) + "\n###End Hypothesis" for label in data['candidate_labels']]
 
 def few_shot_classification(data):
     sequences = create_sequences(data)
     print(sequences)
-    results = onnx_few_shot_model(sequences)
+    # results = onnx_few_shot_model(sequences)
     probs = onnx_few_shot_model.predict_proba(sequences)
-    print(results)
-    print(probs)
-    response_string = json.dumps(results.tolist())
+    response_string = json.dumps(probs.tolist())
     return response_string
 
 gradio_interface = gradio.Interface(