Have multiple options

EleutherAI · Sep 15, 2023 · f6da96b · f6da96b
1 parent feda73f
commit f6da96b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/.gitignore b/.gitignore
@@ -3,7 +3,7 @@
 
 
 !datasets/eval/Pythia_70m_Deduped_Low_Perplexity_Labeling_Formatted.csv
-
+*.pt
 *.zip
 .vscode
 ### Data ###

diff --git a/inference_outputs.py b/inference_outputs.py
@@ -146,8 +146,16 @@ def run_model_inferences(split_name: str, run_id: str, dataset: str, batch_size:
 
 
 
- results = p.map(parse_attn, [t.detach().cpu() for t in outputs.attentions])
- print(results)
+ # results = p.map(parse_attn, [t.detach().cpu() for t in outputs.attentions])
+ # print(results)
+
+ # attentions_table = {}
+ for i in tqdm(range(len(batch[0]))):
+ current_example_id = batch[0][i]
+ current_example_attentions = torch.stack(outputs.attentions)[:, i, :]
+ # attentions_table[current_example_id] = current_example_attentions
+ torch.save(current_example_attentions, f"datasets/{run_id}/{dataset}_attentions_{current_example_id}.pt")
+ # print(current_example_attentions.shape)
 
  # inference_logs = pd.DataFrame({
  # "Loss": outputs.loss.detach().cpu().tolist(),

diff --git a/inference_sync.py b/inference_sync.py
@@ -138,7 +138,7 @@ def get_batch_size(model_name: str) -> int:
  "2.8b": 128,
  # Large
  "6.9b": 64,
- "12b": 64,
+ "12b": 32,
  }
  model_size = ".".join(model_name.split(".")[1:])
  return size_batch_map[model_size]