jit test...

interfaces/cap.py

@@ -85,18 +85,28 @@ def build_huggingface_path(language: str, domain: str):
 
 def predict(text, model_id, tokenizer_id):
     device = torch.device("cpu")
-    model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN).to(device)
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
 
-    inputs = tokenizer(text,
-                       max_length=256,
-                       truncation=True,
-                       padding="do_not_pad",
-                       return_tensors="pt").to(device)
+    # Load JIT-traced model
+    jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
+    model = torch.jit.load(jit_model_path).to(device)
     model.eval()
 
+    # Load tokenizer (still regular HF)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+
+    # Tokenize input
+    inputs = tokenizer(
+        text,
+        max_length=256,
+        truncation=True,
+        padding="do_not_pad",
+        return_tensors="pt"
+    )
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
     with torch.no_grad():
-        logits = model(**inputs).logits
+        logits = model(inputs["input_ids"], inputs["attention_mask"]).logits
+        
     release_model(model, model_id)
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()