Experimenting with few-shot classification.

app.py

@@ -74,7 +74,7 @@ few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-ss
 candidate_labels = ["supported", "refuted"]
 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="This hypothesis is {}")
+train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="Based on the Given passage, the hypothesis is {}.")
 args = TrainingArguments(
     batch_size=32,
     num_epochs=1
@@ -110,13 +110,13 @@ def zero_shot_classification(data):
     return response_string
 
 def create_sequences(data):
-    return ['###Given:\n' + data['sequence'] + '\n###Hypothesis:\n' + data['hypothesis_template'].format(label) for label in data['candidate_labels']]
+    return ['###Given:\n' + data['sequence'] + '\n###End Given\n###Hypothesis:\n' + data['hypothesis_template'].format(label) for label in data['candidate_labels'] + "\n###End Hypothesis" + "\nBased on the Given passage, the Hypothesis is {}"]
 
 def few_shot_classification(data):
     sequences = create_sequences(data)
     print(sequences)
     results = onnx_few_shot_model(sequences)
-    probs = onnx_few_shot_model.predict_proba(data['sequence'])
+    probs = onnx_few_shot_model.predict_proba(sequences)
     print(results)
     print(probs)
     response_string = json.dumps(results.tolist())