JIT...

utils.py

@@ -25,8 +25,7 @@ from interfaces.ontolisst import build_huggingface_path as hf_ontolisst_path
 
 from huggingface_hub import scan_cache_dir
 
-MODELS_PRELOADED = []
-TOKENIZERS_PRELOADED = []
+JIT_DIR = "/data/jit_models"
 
 HF_TOKEN = os.environ["hf_read"]
 
@@ -55,12 +54,53 @@ for domain in domains_illframes.values():
 
 tokenizers = ["xlm-roberta-large"]
 
-def download_hf_models():
+def download_hf_models(models=[], tokenizers=[], hf_token=None):
+    # Ensure the JIT model directory exists
+    os.makedirs(JIT_DIR, exist_ok=True)
+
     for model_id in models:
-        MODELS_PRELOADED[model_id] = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
-    for tokenizer_id in tokenizers:
-        TOKENIZERS_PRELOADED[tokenizer_id] = AutoTokenizer.from_pretrained(tokenizer_id)
+        print(f"Downloading + JIT tracing model: {model_id}")
         
+        # Load model and tokenizer
+        model = AutoModelForSequenceClassification.from_pretrained(
+            model_id,
+            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
+        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
+        )
         
 def df_h():
     result = subprocess.run(["df", "-H"], capture_output=True, text=True)