update endpoint helper files

infer_local.py

@@ -286,13 +286,19 @@ def main():
     parser.add_argument("--base-model", type=str, default="Qwen/Qwen2.5-Coder-0.5B-Instruct")
     parser.add_argument("--prompt", type=str, required=True)
     parser.add_argument("--max-new-tokens", type=int, default=320)
-    parser.add_argument("--temperature", type=float, default=0.25)
-    parser.add_argument("--top-p", type=float, default=0.9)
-    parser.add_argument("--do-sample", action="store_true")
-    args = parser.parse_args()
-
-    if not os.path.exists(args.model_path):
-        raise FileNotFoundError(
+    parser.add_argument("--temperature", type=float, default=0.25)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--do-sample", action="store_true")
+    parser.add_argument(
+        "--allow-downloads",
+        action="store_true",
+        help="Allow Transformers to download missing model files from Hugging Face.",
+    )
+    args = parser.parse_args()
+    local_files_only = not args.allow_downloads
+
+    if not os.path.exists(args.model_path):
+        raise FileNotFoundError(
             f"Model path not found: {args.model_path}. Train first using run_pipeline.py."
         )
 
@@ -301,20 +307,28 @@ def main():
     full_model_weights_present = has_full_model_weights(args.model_path)
 
     if os.path.exists(adapter_config_path) and adapter_weights_present:
-        peft_config = PeftConfig.from_pretrained(args.model_path)
-        base_model_name = peft_config.base_model_name_or_path or args.base_model
-        tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-        base_model = AutoModelForCausalLM.from_pretrained(
-            base_model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        )
-        model = PeftModel.from_pretrained(base_model, args.model_path)
-    elif full_model_weights_present and not os.path.exists(adapter_config_path):
-        tokenizer = AutoTokenizer.from_pretrained(args.model_path)
-        model = AutoModelForCausalLM.from_pretrained(
-            args.model_path,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        )
+        peft_config = PeftConfig.from_pretrained(args.model_path)
+        base_model_name = peft_config.base_model_name_or_path or args.base_model
+        tokenizer = AutoTokenizer.from_pretrained(
+            base_model_name,
+            local_files_only=local_files_only,
+        )
+        base_model = AutoModelForCausalLM.from_pretrained(
+            base_model_name,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            local_files_only=local_files_only,
+        )
+        model = PeftModel.from_pretrained(base_model, args.model_path)
+    elif full_model_weights_present and not os.path.exists(adapter_config_path):
+        tokenizer = AutoTokenizer.from_pretrained(
+            args.model_path,
+            local_files_only=local_files_only,
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_path,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            local_files_only=local_files_only,
+        )
     else:
         # Graceful fallback when local model folder has config/tokenizer but no weight files.
         fallback_base = args.base_model
@@ -343,11 +357,15 @@ def main():
                 ),
                 file=sys.stderr,
             )
-        tokenizer = AutoTokenizer.from_pretrained(fallback_base)
-        model = AutoModelForCausalLM.from_pretrained(
-            fallback_base,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            fallback_base,
+            local_files_only=local_files_only,
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            fallback_base,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            local_files_only=local_files_only,
+        )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     model.eval()