Experimenting with few-shot classification.

app.py

@@ -71,16 +71,16 @@ 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 = ["true", "false"]
+candidate_labels = ["correct", "wrong"]
 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 statement is {}.")
+train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="This conclusion is {}.")
 args = TrainingArguments(
     batch_size=32,
     num_epochs=1
 )
 trainer = Trainer(
-    model=few_shot_model,
+    model=ort_model,
     args=args,
     train_dataset=train_dataset,
     eval_dataset=reference_dataset["test"]