prepend NaN values to logprobs and clean up

scripts/inference.py

@@ -4,6 +4,7 @@
 import pandas as pd
 import torch
 from datasets import Dataset, load_dataset
+from jaxtyping import Int
 from tqdm.auto import tqdm
 from transformers import AutoModelForCausalLM
 
@@ -12,15 +13,33 @@
 torch.set_grad_enabled(False)
 
 
-def main(model_name, batch_size, dataset_name, token):
+def main(model_name: str, batch_size: Int, dataset_name: str, token: str):
+    """
+    Outputs the log probabilities of the next token for each token in the validation dataset.
+    And uploads the resulting dataset to huggingface.
+    Args:
+    - model_name: The name of the model to use for inference
+    - batch_size: The batch size for processing. 80 worked well in CPU.
+    - dataset_name: The name of the dataset from which validation set will be loaded
+    - token: Hugging Face API token
+    """
     val_ds = load_validation_dataset(dataset_name)
 
     # model accepts 2D tensors (batch_size, seq_len)
     val_sequences = torch.tensor([s["tokens"] for s in val_ds])
     model = AutoModelForCausalLM.from_pretrained(model_name)
-    logprobs, next_logprobs = get_all_and_next_logprobs(model, val_sequences)
 
-    df_dataset = pd.DataFrame({"logprobs": next_logprobs.tolist()})
+    accumulated_logprobs = torch.tensor([], dtype=torch.float32)
+
+    for i in tqdm(range(0, len(val_sequences), batch_size)):
+        batch_sequences = val_sequences[i : i + batch_size]
+        _, next_logprobs = get_all_and_next_logprobs(model, val_sequences)
+        accumulated_logprobs = torch.cat((accumulated_logprobs, next_logprobs), dim=0)
+
+    nan_tensor = torch.full((accumulated_logprobs.size(0), 1), float("nan"))
+    extended_next_logprobs = torch.cat([nan_tensor, next_logprobs], dim=1)  # 513 tokens
+
+    df_dataset = pd.DataFrame({"logprobs": extended_next_logprobs.tolist()})
     hf_dataset = Dataset.from_pandas(df_dataset)
 
     # change the repo_id to your hf username