kovacsvi commited on
Commit
fb1a253
·
1 Parent(s): 0b09517

JIT tracing

Browse files
interfaces/cap.py CHANGED
@@ -106,7 +106,7 @@ def predict(text, model_id, tokenizer_id):
106
 
107
  with torch.no_grad():
108
  output = model(inputs["input_ids"], inputs["attention_mask"])
109
- print(output)
110
  logits = output["logits"]
111
 
112
  release_model(model, model_id)
 
106
 
107
  with torch.no_grad():
108
  output = model(inputs["input_ids"], inputs["attention_mask"])
109
+ print(output) # debug
110
  logits = output["logits"]
111
 
112
  release_model(model, model_id)
interfaces/cap_media_demo.py CHANGED
@@ -35,18 +35,30 @@ def build_huggingface_path(language: str, domain: str):
35
 
36
  def predict(text, model_id, tokenizer_id):
37
  device = torch.device("cpu")
38
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
39
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
40
 
41
- inputs = tokenizer(text,
42
- max_length=256,
43
- truncation=True,
44
- padding="do_not_pad",
45
- return_tensors="pt").to(device)
46
  model.eval()
47
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48
  with torch.no_grad():
49
- logits = model(**inputs).logits
 
 
 
50
  release_model(model, model_id)
51
 
52
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
35
 
36
  def predict(text, model_id, tokenizer_id):
37
  device = torch.device("cpu")
 
 
38
 
39
+ # Load JIT-traced model
40
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
41
+ model = torch.jit.load(jit_model_path).to(device)
 
 
42
  model.eval()
43
 
44
+ # Load tokenizer (still regular HF)
45
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
46
+
47
+ # Tokenize input
48
+ inputs = tokenizer(
49
+ text,
50
+ max_length=256,
51
+ truncation=True,
52
+ padding="do_not_pad",
53
+ return_tensors="pt"
54
+ )
55
+ inputs = {k: v.to(device) for k, v in inputs.items()}
56
+
57
  with torch.no_grad():
58
+ output = model(inputs["input_ids"], inputs["attention_mask"])
59
+ print(output) # debug
60
+ logits = output["logits"]
61
+
62
  release_model(model, model_id)
63
 
64
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/cap_minor.py CHANGED
@@ -67,18 +67,30 @@ def build_huggingface_path(language: str, domain: str):
67
 
68
  def predict(text, model_id, tokenizer_id):
69
  device = torch.device("cpu")
70
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
71
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
72
 
73
- inputs = tokenizer(text,
74
- max_length=256,
75
- truncation=True,
76
- padding="do_not_pad",
77
- return_tensors="pt").to(device)
78
  model.eval()
79
 
 
 
 
 
 
 
 
 
 
 
 
 
 
80
  with torch.no_grad():
81
- logits = model(**inputs).logits
 
 
 
82
  release_model(model, model_id)
83
 
84
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
67
 
68
  def predict(text, model_id, tokenizer_id):
69
  device = torch.device("cpu")
 
 
70
 
71
+ # Load JIT-traced model
72
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
73
+ model = torch.jit.load(jit_model_path).to(device)
 
 
74
  model.eval()
75
 
76
+ # Load tokenizer (still regular HF)
77
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
78
+
79
+ # Tokenize input
80
+ inputs = tokenizer(
81
+ text,
82
+ max_length=256,
83
+ truncation=True,
84
+ padding="do_not_pad",
85
+ return_tensors="pt"
86
+ )
87
+ inputs = {k: v.to(device) for k, v in inputs.items()}
88
+
89
  with torch.no_grad():
90
+ output = model(inputs["input_ids"], inputs["attention_mask"])
91
+ print(output) # debug
92
+ logits = output["logits"]
93
+
94
  release_model(model, model_id)
95
 
96
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/cap_minor_media.py CHANGED
@@ -150,18 +150,30 @@ def predict(text, major_model_id, minor_model_id, tokenizer_id, HF_TOKEN=None):
150
 
151
  def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
152
  device = torch.device("cpu")
153
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN).to(device)
154
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
155
 
156
- inputs = tokenizer(text,
157
- max_length=256,
158
- truncation=True,
159
- padding="do_not_pad",
160
- return_tensors="pt").to(device)
161
  model.eval()
162
 
 
 
 
 
 
 
 
 
 
 
 
 
 
163
  with torch.no_grad():
164
- logits = model(**inputs).logits
 
 
 
165
  release_model(model, model_id)
166
 
167
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
150
 
151
  def predict_flat(text, model_id, tokenizer_id, HF_TOKEN=None):
152
  device = torch.device("cpu")
 
 
153
 
154
+ # Load JIT-traced model
155
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
156
+ model = torch.jit.load(jit_model_path).to(device)
 
 
157
  model.eval()
158
 
159
+ # Load tokenizer (still regular HF)
160
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
161
+
162
+ # Tokenize input
163
+ inputs = tokenizer(
164
+ text,
165
+ max_length=256,
166
+ truncation=True,
167
+ padding="do_not_pad",
168
+ return_tensors="pt"
169
+ )
170
+ inputs = {k: v.to(device) for k, v in inputs.items()}
171
+
172
  with torch.no_grad():
173
+ output = model(inputs["input_ids"], inputs["attention_mask"])
174
+ print(output) # debug
175
+ logits = output["logits"]
176
+
177
  release_model(model, model_id)
178
 
179
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/emotion.py CHANGED
@@ -27,19 +27,30 @@ def build_huggingface_path(language: str):
27
 
28
  def predict(text, model_id, tokenizer_id):
29
  device = torch.device("cpu")
30
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
31
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
32
- model.to(device)
33
 
34
- inputs = tokenizer(text,
35
- max_length=512,
36
- truncation=True,
37
- padding="do_not_pad",
38
- return_tensors="pt").to(device)
39
  model.eval()
40
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
  with torch.no_grad():
42
- logits = model(**inputs).logits
 
 
 
43
  release_model(model, model_id)
44
 
45
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
27
 
28
  def predict(text, model_id, tokenizer_id):
29
  device = torch.device("cpu")
 
 
 
30
 
31
+ # Load JIT-traced model
32
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
33
+ model = torch.jit.load(jit_model_path).to(device)
 
 
34
  model.eval()
35
 
36
+ # Load tokenizer (still regular HF)
37
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
38
+
39
+ # Tokenize input
40
+ inputs = tokenizer(
41
+ text,
42
+ max_length=256,
43
+ truncation=True,
44
+ padding="do_not_pad",
45
+ return_tensors="pt"
46
+ )
47
+ inputs = {k: v.to(device) for k, v in inputs.items()}
48
+
49
  with torch.no_grad():
50
+ output = model(inputs["input_ids"], inputs["attention_mask"])
51
+ print(output) # debug
52
+ logits = output["logits"]
53
+
54
  release_model(model, model_id)
55
 
56
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/emotion9.py CHANGED
@@ -26,18 +26,30 @@ def build_huggingface_path(language: str):
26
 
27
  def predict(text, model_id, tokenizer_id):
28
  device = torch.device("cpu")
29
- model = AutoModelForSequenceClassification.from_pretrained(model_id, token=HF_TOKEN)
30
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
31
 
32
- inputs = tokenizer(text,
33
- max_length=512,
34
- truncation=True,
35
- padding="do_not_pad",
36
- return_tensors="pt").to(device)
37
  model.eval()
38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  with torch.no_grad():
40
- logits = model(**inputs).logits
 
 
 
41
  release_model(model, model_id)
42
 
43
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
26
 
27
  def predict(text, model_id, tokenizer_id):
28
  device = torch.device("cpu")
 
 
29
 
30
+ # Load JIT-traced model
31
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
32
+ model = torch.jit.load(jit_model_path).to(device)
 
 
33
  model.eval()
34
 
35
+ # Load tokenizer (still regular HF)
36
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
37
+
38
+ # Tokenize input
39
+ inputs = tokenizer(
40
+ text,
41
+ max_length=256,
42
+ truncation=True,
43
+ padding="do_not_pad",
44
+ return_tensors="pt"
45
+ )
46
+ inputs = {k: v.to(device) for k, v in inputs.items()}
47
+
48
  with torch.no_grad():
49
+ output = model(inputs["input_ids"], inputs["attention_mask"])
50
+ print(output) # debug
51
+ logits = output["logits"]
52
+
53
  release_model(model, model_id)
54
 
55
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/illframes.py CHANGED
@@ -58,28 +58,30 @@ def build_huggingface_path(domain: str):
58
 
59
  def predict(text, model_id, tokenizer_id, label_names):
60
  device = torch.device("cpu")
61
- try:
62
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
63
- except:
64
- disk_space = get_disk_space('/data/')
65
- print("Disk Space Error:")
66
- for key, value in disk_space.items():
67
- print(f"{key}: {value}")
68
 
69
- shutil.rmtree("/data")
70
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN, force_download=True)
71
-
 
 
 
72
  tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
73
 
74
- inputs = tokenizer(text,
75
- max_length=256,
76
- truncation=True,
77
- padding="do_not_pad",
78
- return_tensors="pt").to(device)
79
- model.eval()
 
 
 
80
 
81
  with torch.no_grad():
82
- logits = model(**inputs).logits
 
 
 
83
  release_model(model, model_id)
84
 
85
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
58
 
59
  def predict(text, model_id, tokenizer_id, label_names):
60
  device = torch.device("cpu")
 
 
 
 
 
 
 
61
 
62
+ # Load JIT-traced model
63
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
64
+ model = torch.jit.load(jit_model_path).to(device)
65
+ model.eval()
66
+
67
+ # Load tokenizer (still regular HF)
68
  tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
69
 
70
+ # Tokenize input
71
+ inputs = tokenizer(
72
+ text,
73
+ max_length=256,
74
+ truncation=True,
75
+ padding="do_not_pad",
76
+ return_tensors="pt"
77
+ )
78
+ inputs = {k: v.to(device) for k, v in inputs.items()}
79
 
80
  with torch.no_grad():
81
+ output = model(inputs["input_ids"], inputs["attention_mask"])
82
+ print(output) # debug
83
+ logits = output["logits"]
84
+
85
  release_model(model, model_id)
86
 
87
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/manifesto.py CHANGED
@@ -26,18 +26,30 @@ def build_huggingface_path(language: str):
26
 
27
  def predict(text, model_id, tokenizer_id):
28
  device = torch.device("cpu")
29
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
30
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
31
 
32
- inputs = tokenizer(text,
33
- max_length=256,
34
- truncation=True,
35
- padding="do_not_pad",
36
- return_tensors="pt").to(device)
37
  model.eval()
38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  with torch.no_grad():
40
- logits = model(**inputs).logits
 
 
 
41
  release_model(model, model_id)
42
 
43
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
26
 
27
  def predict(text, model_id, tokenizer_id):
28
  device = torch.device("cpu")
 
 
29
 
30
+ # Load JIT-traced model
31
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
32
+ model = torch.jit.load(jit_model_path).to(device)
 
 
33
  model.eval()
34
 
35
+ # Load tokenizer (still regular HF)
36
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
37
+
38
+ # Tokenize input
39
+ inputs = tokenizer(
40
+ text,
41
+ max_length=256,
42
+ truncation=True,
43
+ padding="do_not_pad",
44
+ return_tensors="pt"
45
+ )
46
+ inputs = {k: v.to(device) for k, v in inputs.items()}
47
+
48
  with torch.no_grad():
49
+ output = model(inputs["input_ids"], inputs["attention_mask"])
50
+ print(output) # debug
51
+ logits = output["logits"]
52
+
53
  release_model(model, model_id)
54
 
55
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/ontolisst.py CHANGED
@@ -44,29 +44,30 @@ def build_huggingface_path(language: str):
44
 
45
  def predict(text, model_id, tokenizer_id):
46
  device = torch.device("cpu")
47
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
48
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
49
-
50
- # --- DEBUG ---
51
-
52
- disk_space = get_disk_space('/data/')
53
- print("Disk Space Info:")
54
- for key, value in disk_space.items():
55
- print(f"{key}: {value}")
56
-
57
- # ---
58
-
59
- model.to(device)
60
 
61
- inputs = tokenizer(text,
62
- max_length=256,
63
- truncation=True,
64
- padding="do_not_pad",
65
- return_tensors="pt").to(device)
66
  model.eval()
67
 
 
 
 
 
 
 
 
 
 
 
 
 
 
68
  with torch.no_grad():
69
- logits = model(**inputs).logits
 
 
 
70
  release_model(model, model_id)
71
 
72
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
44
 
45
  def predict(text, model_id, tokenizer_id):
46
  device = torch.device("cpu")
 
 
 
 
 
 
 
 
 
 
 
 
 
47
 
48
+ # Load JIT-traced model
49
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
50
+ model = torch.jit.load(jit_model_path).to(device)
 
 
51
  model.eval()
52
 
53
+ # Load tokenizer (still regular HF)
54
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
55
+
56
+ # Tokenize input
57
+ inputs = tokenizer(
58
+ text,
59
+ max_length=256,
60
+ truncation=True,
61
+ padding="do_not_pad",
62
+ return_tensors="pt"
63
+ )
64
+ inputs = {k: v.to(device) for k, v in inputs.items()}
65
+
66
  with torch.no_grad():
67
+ output = model(inputs["input_ids"], inputs["attention_mask"])
68
+ print(output) # debug
69
+ logits = output["logits"]
70
+
71
  release_model(model, model_id)
72
 
73
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
interfaces/sentiment.py CHANGED
@@ -30,19 +30,30 @@ def build_huggingface_path(language: str):
30
 
31
  def predict(text, model_id, tokenizer_id):
32
  device = torch.device("cpu")
33
- model = AutoModelForSequenceClassification.from_pretrained(model_id, device_map="auto", token=HF_TOKEN)
34
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
35
- model.to(device)
36
 
37
- inputs = tokenizer(text,
38
- max_length=256,
39
- truncation=True,
40
- padding="do_not_pad",
41
- return_tensors="pt").to(device)
42
  model.eval()
43
 
 
 
 
 
 
 
 
 
 
 
 
 
 
44
  with torch.no_grad():
45
- logits = model(**inputs).logits
 
 
 
46
  release_model(model, model_id)
47
 
48
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
 
30
 
31
  def predict(text, model_id, tokenizer_id):
32
  device = torch.device("cpu")
 
 
 
33
 
34
+ # Load JIT-traced model
35
+ jit_model_path = f"/data/jit_models/{model_id.replace('/', '_')}.pt"
36
+ model = torch.jit.load(jit_model_path).to(device)
 
 
37
  model.eval()
38
 
39
+ # Load tokenizer (still regular HF)
40
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
41
+
42
+ # Tokenize input
43
+ inputs = tokenizer(
44
+ text,
45
+ max_length=256,
46
+ truncation=True,
47
+ padding="do_not_pad",
48
+ return_tensors="pt"
49
+ )
50
+ inputs = {k: v.to(device) for k, v in inputs.items()}
51
+
52
  with torch.no_grad():
53
+ output = model(inputs["input_ids"], inputs["attention_mask"])
54
+ print(output) # debug
55
+ logits = output["logits"]
56
+
57
  release_model(model, model_id)
58
 
59
  probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
utils.py CHANGED
@@ -75,7 +75,7 @@ def download_hf_models():
75
  safe_model_name = model_id.replace("/", "_")
76
  traced_model_path = os.path.join(JIT_DIR, f"{safe_model_name}.pt")
77
 
78
- if os.path.exists(traced_model_path) or "pooled-cap" not in model_id:
79
  print(f"⏩ Skipping JIT — already exists: {traced_model_path}")
80
  else:
81
  print(f"⚙️ Tracing and saving: {traced_model_path}")
 
75
  safe_model_name = model_id.replace("/", "_")
76
  traced_model_path = os.path.join(JIT_DIR, f"{safe_model_name}.pt")
77
 
78
+ if os.path.exists(traced_model_path):
79
  print(f"⏩ Skipping JIT — already exists: {traced_model_path}")
80
  else:
81
  print(f"⚙️ Tracing and saving: {traced_model_path}")