Experimenting with few-shot classification.

app.py

@@ -57,22 +57,17 @@ classifier = pipeline(task="zero-shot-classification", model=model, tokenizer=to
 
 few_shot_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5', model_max_length=512)
 ort_model = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5', file_name="onnx/model.onnx")
-few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
+few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english", multi_target_strategy="multi-output")
 
-
-
-candidate_labels = ["true", "false"]
+# Train few_shot_model
+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 statement is {}")
-
-
-
+train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="This hypothesis is {}")
 args = TrainingArguments(
     batch_size=32,
     num_epochs=1
 )
-
 trainer = Trainer(
     model=few_shot_model,
     args=args,
@@ -81,8 +76,8 @@ trainer = Trainer(
 )
 trainer.train()
 
-metrics = trainer.evaluate()
-print(metrics)
+# metrics = trainer.evaluate()
+# print(metrics)
 
 onnx_few_shot_model = OnnxSetFitModel(ort_model, few_shot_tokenizer, few_shot_model.model_head)