remove input pos keyword in generate.py

pytorch · Sep 24, 2024 · 937e7ed · 937e7ed
1 parent a70d7b5
commit 937e7ed
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 472 deletions.
diff --git a/torchchat/cli/download.py b/torchchat/cli/download.py
@@ -30,7 +30,7 @@ def _download_hf_snapshot(
     try:
         snapshot_download(
             model_config.distribution_path,
-            cache_dir=artifact_dir,
+            local_dir=artifact_dir,
             local_dir_use_symlinks=False,
             token=hf_token,
             ignore_patterns=None if "llava" in model_config.name else "*safetensors*",

diff --git a/torchchat/generate.py b/torchchat/generate.py
@@ -374,7 +374,7 @@ def prefill(
             logits = model(x)
         else:
             # input_pos: [B, S]
-            logits = model(x, input_pos=input_pos)
+            logits = model(x, input_pos)
             # print(f"logits {logits.shape}")
 
         # print(f"x: {x},\n  input_pos: {input_pos}\n")
@@ -398,7 +398,7 @@ def decode_one_token(
             else:
                 logits = model(x)
         else:
-            logits = model(x, input_pos=input_pos)
+            logits = model(x, input_pos)
         # print(f"x: {x},\n  input_pos: {input_pos}\n")
         return self.sample(logits, need_probs=need_probs, **sampling_kwargs)