Experimenting with few-shot classification.

app.py

@@ -59,15 +59,14 @@ 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")
 
-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=candidate_labels, sample_size=8, template="This statement is {}")
+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 {}")
+
+
 
 args = TrainingArguments(
     batch_size=32,
@@ -76,8 +75,8 @@ args = TrainingArguments(
 
 trainer = Trainer(
     model=few_shot_model,
-    args=args,
-    train_dataset=synthetic_dataset
+    train_dataset=train_dataset,
+    eval_dataset=reference_dataset["test"]
 )
 trainer.train()