implementing knn-based influential dataset approach

src/small_model_training/config.json

@@ -1,7 +1,7 @@
 {
  "model_name_or_path": "bert-base-uncased",
  "tokenizer_name": "distilbert-base-uncased",
- "output_dir":"my_awesome_model",
+ "output_dir":"imdb_model",
  "learning_rate":2e-5,
  "per_device_train_batch_size":16,
  "per_device_eval_batch_size":16,

src/small_model_training/create_inf_subset.py

@@ -0,0 +1,64 @@
+import torch.nn
+import numpy as np
+from transformers import AutoTokenizer
+from transformers import AutoModelForSequenceClassification
+from transformers import DataCollatorWithPadding
+from transformers import Trainer
+from datasets import load_dataset
+import evaluate
+
+def compute_metrics(eval_pred):
+ predictions, labels = eval_pred
+ predictions = np.argmax(predictions, axis=1)
+ return accuracy.compute(predictions=predictions, references=labels)
+
+accuracy = evaluate.load("accuracy")
+
+model = AutoModelForSequenceClassification.from_pretrained("./imdb_model")
+
+imdb_train = load_dataset("imdb", split="train[:10%]")
+imdb_test = load_dataset("imdb", split="test[:1%]")
+
+tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+
+
+def preprocess_function(examples):
+ return tokenizer(examples["text"], truncation=True)
+
+
+tokenized_imdb_train = imdb_train.map(preprocess_function, batched=True)
+tokenized_imdb_test = imdb_test.map(preprocess_function, batched=True)
+
+data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+
+trainer = Trainer(
+ model=model,
+ train_dataset=tokenized_imdb_train,
+ eval_dataset=tokenized_imdb_test,
+ tokenizer=tokenizer,
+ data_collator=data_collator,
+ compute_metrics=compute_metrics,
+)
+
+outputs = trainer.predict(tokenized_imdb_test)
+
+logits = outputs[0]
+
+logits = torch.from_numpy(logits)
+
+scores = torch.nn.functional.softmax(logits, dim=-1)
+
+first_values = scores[:, 0]
+second_values = scores[:, 1]
+
+distance = first_values-second_values
+
+dist = torch.abs(distance)
+knn_values, knn_indices = dist.topk(5, largest=False)
+
+fewshot_examples = []
+
+for elem in knn_indices:
+ fewshot_examples.append(imdb_test[elem.item()]["text"])
+
+print(fewshot_examples)

src/small_model_training/text_classification.py

@@ -26,6 +26,8 @@ def preprocess_function(examples):
 
  tokenized_imdb = dataset.map(preprocess_function, batched=True)
 
+ print("done tokenizing")
+
  data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
  id2label = {0: "NEGATIVE", 1: "POSITIVE"}
@@ -46,6 +48,9 @@ def preprocess_function(examples):
  )
 
  trainer.train()
+ trainer.evaluate()
+
+ model.save_pretrained("./imdb_model")
 
 
  # TO-DO: calculate influential dataset