jit tracing fix

utils.py

@@ -69,40 +69,37 @@ def download_hf_models():
             token=HF_TOKEN,
             device_map="auto"
         )
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
-            token=HF_TOKEN
-        )
-
-        model.eval()
-
-        # Dummy input for tracing
-        dummy_input = tokenizer(
-            "Hello, world!",
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=256
-        )
-
-        # JIT trace
-        traced_model = torch.jit.trace(
-            model,
-            (dummy_input["input_ids"], dummy_input["attention_mask"])
-        )
-
-        # Save traced model
+        tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-large")
+        
         safe_model_name = model_id.replace("/", "_")
         traced_model_path = os.path.join(JIT_DIR, f"{safe_model_name}.pt")
-        traced_model.save(traced_model_path)
-        print(f"✔️ Saved JIT model to: {traced_model_path}")
-
-    for tokenizer_id in tokenizers:
-        print(f"Downloading tokenizer: {tokenizer_id}")
-        AutoTokenizer.from_pretrained(
-            tokenizer_id,
-            token=HF_TOKEN
-        )
+        
+        if os.path.exists(traced_model_path):
+            print(f"⏩ Skipping JIT — already exists: {traced_model_path}")
+        else:
+            print(f"⚙️  Tracing and saving: {traced_model_path}")
+
+            model.eval()
+
+            # Dummy input for tracing
+            dummy_input = tokenizer(
+                "Hello, world!",
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=256
+            )
+
+            # JIT trace
+            traced_model = torch.jit.trace(
+                model,
+                (dummy_input["input_ids"], dummy_input["attention_mask"]),
+                strict=False
+            )
+
+            # Save traced model
+            traced_model.save(traced_model_path)
+            print(f"✔️ Saved JIT model to: {traced_model_path}")
         
 def df_h():
     result = subprocess.run(["df", "-H"], capture_output=True, text=True)