Experimenting with few-shot classification.

app.py

@@ -59,24 +59,25 @@ few_shot_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-small-en-v1.5', mod
 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", multi_target_strategy="multi-output")
 
-test_dataset = load_dataset("dair-ai/emotion", "split", split="test")
-print(test_dataset)
-classes = test_dataset.features["label"].names
-print(classes)
-train_dataset = get_templated_dataset(candidate_labels=classes)
-print(train_dataset)
-print(train_dataset[0])
+def get_templated_dataset(candidate_labels, sample_size, template="This sentence is {}"):
+    examples = []
+    for label in candidate_labels:
+        for _ in range(sample_size):
+            examples.append({"text": template.format(label), "label": label})
+    return examples
+
+candidate_labels = ["true", "false"]
+synthetic_dataset = get_templated_dataset(candidate_labels, sample_size=8)
 
 args = TrainingArguments(
     batch_size=32,
     num_epochs=1
 )
 
-trainer = Trainer(
+trainer = SetFitTrainer(
     model=few_shot_model,
     args=args,
-    train_dataset=train_dataset,
-    eval_dataset=test_dataset
+    train_dataset=train_dataset
 )
 trainer.train()