memory usage improved

delphi-suite · Feb 21, 2024 · 51b1c57 · 51b1c57
1 parent 437c250
commit 51b1c57
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 11 deletions.
diff --git a/scripts/inference.py b/scripts/inference.py
@@ -51,16 +51,7 @@ def main(
         logprobs_tensor = get_all_and_next_logprobs(model, batch_sequences_tensor)[1]
         logprobs[i:batch_end, 1:] = logprobs_tensor.cpu().numpy()
 
-        # logprobs_list.append(next_logprobs)
-
-    # accumulated_logprobs = torch.cat(logprobs_list, dim=0)
-
-    # nan_tensor = torch.full((accumulated_logprobs.size(0), 1), float("nan"))
-    # extended_next_logprobs = torch.cat(
-    #     [nan_tensor, accumulated_logprobs], dim=1
-    # )  # 513 tokens
-
-    df_dataset = pd.DataFrame({"logprobs": logprobs.tolist()})
+    df_dataset = pd.DataFrame({"logprobs": [row for row in logprobs]})
     hf_dataset = Dataset.from_pandas(df_dataset)
 
     # change the repo_id to your hf username in generate_logprobs.sh

diff --git a/tests/scripts/functional_test_generate_logprobs.sh b/tests/scripts/functional_test_generate_logprobs.sh
@@ -5,7 +5,7 @@
 BATCH_SIZE=80
 DATASET_NAME="delphi-suite/tinystories-v2-clean-tokenized"
 USERNAME="transcendingvictor"  # Your Hugging Face username
-TOKEN="hf_aaaaaaaaaaaaaaaaaaaaaaaaaaaaa"  # Your Hugging Face API token
+TOKEN="hf_aaaaaaaaaaaaaaaaaaaaaaaaa"  # Your Hugging Face API token
 
 # List of models
 declare -a MODEL_NAMES=("delphi-suite/delphi-llama2-100k"