Lord-Raven commited on
Commit
07c25e9
·
1 Parent(s): 30dda8a

Experimenting with few-shot classification.

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -71,16 +71,16 @@ ort_model = ORTModelForFeatureExtraction.from_pretrained('BAAI/bge-small-en-v1.5
71
  few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
72
 
73
  # Train few_shot_model
74
- candidate_labels = ["true", "false"]
75
  reference_dataset = load_dataset("emotion")
76
  dummy_dataset = Dataset.from_dict({})
77
- train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="This statement is {}.")
78
  args = TrainingArguments(
79
  batch_size=32,
80
  num_epochs=1
81
  )
82
  trainer = Trainer(
83
- model=few_shot_model,
84
  args=args,
85
  train_dataset=train_dataset,
86
  eval_dataset=reference_dataset["test"]
 
71
  few_shot_model = SetFitModel.from_pretrained("moshew/bge-small-en-v1.5_setfit-sst2-english")
72
 
73
  # Train few_shot_model
74
+ candidate_labels = ["correct", "wrong"]
75
  reference_dataset = load_dataset("emotion")
76
  dummy_dataset = Dataset.from_dict({})
77
+ train_dataset = get_templated_dataset(dummy_dataset, candidate_labels=candidate_labels, sample_size=8, template="This conclusion is {}.")
78
  args = TrainingArguments(
79
  batch_size=32,
80
  num_epochs=1
81
  )
82
  trainer = Trainer(
83
+ model=ort_model,
84
  args=args,
85
  train_dataset=train_dataset,
86
  eval_dataset=reference_dataset["test"]