vickeee465 commited on
Commit
e1df718
·
1 Parent(s): 093e523

finding bottlenecks pt1

Browse files
Files changed (1) hide show
  1. interfaces/manifesto.py +37 -17
interfaces/manifesto.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import gradio as gr
2
 
3
  import os
@@ -9,6 +10,18 @@ from huggingface_hub import HfApi
9
 
10
  from label_dicts import MANIFESTO_LABEL_NAMES
11
 
 
 
 
 
 
 
 
 
 
 
 
 
12
  HF_TOKEN = os.environ["hf_read"]
13
 
14
  languages = [
@@ -24,23 +37,30 @@ def build_huggingface_path(language: str):
24
 
25
  def predict(text, model_id, tokenizer_id):
26
  device = torch.device("cpu")
27
- model = AutoModelForSequenceClassification.from_pretrained(model_id, token=HF_TOKEN)
28
- tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
29
- model.to(device)
30
-
31
- inputs = tokenizer(text,
32
- max_length=512,
33
- truncation=True,
34
- padding="do_not_pad",
35
- return_tensors="pt").to(device)
36
- model.eval()
37
-
38
- with torch.no_grad():
39
- logits = model(**inputs).logits
40
-
41
- probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
42
- output_pred = {f"[{model.config.id2label[i]}] {MANIFESTO_LABEL_NAMES[int(model.config.id2label[i])]}": probs[i] for i in np.argsort(probs)[::-1]}
43
- output_info = f'<p style="text-align: center; display: block">Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.</p>'
 
 
 
 
 
 
 
44
  return output_pred, output_info
45
 
46
  def predict_cap(text, language):
 
1
+ import time
2
  import gradio as gr
3
 
4
  import os
 
10
 
11
  from label_dicts import MANIFESTO_LABEL_NAMES
12
 
13
+ class RuntimeMeasure:
14
+ def __enter__(self):
15
+ self.start_time = time.time()
16
+ return self
17
+
18
+ def __exit__(self, exc_type, exc_value, traceback):
19
+ end_time = time.time()
20
+ runtime = end_time - self.start_time
21
+ gr.Info(f"Runtime: {runtime} seconds")
22
+ def m(msg):
23
+ return RuntimeMeasure(msg)
24
+
25
  HF_TOKEN = os.environ["hf_read"]
26
 
27
  languages = [
 
37
 
38
  def predict(text, model_id, tokenizer_id):
39
  device = torch.device("cpu")
40
+ with m("Loading model"):
41
+ model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
42
+ with m("Loading tokenizer"):
43
+ tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
44
+
45
+ with m("Tokenizing"):
46
+ inputs = tokenizer(text,
47
+ max_length=512,
48
+ truncation=True,
49
+ padding="do_not_pad",
50
+ return_tensors="pt").to(device)
51
+ with m("model.eval()"):
52
+ model.eval()
53
+
54
+ with m("Inference"):
55
+ with torch.no_grad():
56
+ logits = model(**inputs).logits
57
+
58
+ with m("Softmax"):
59
+ probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
60
+
61
+ with m("Output formatting"):
62
+ output_pred = {f"[{model.config.id2label[i]}] {MANIFESTO_LABEL_NAMES[int(model.config.id2label[i])]}": probs[i] for i in np.argsort(probs)[::-1]}
63
+ output_info = f'<p style="text-align: center; display: block">Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.</p>'
64
  return output_pred, output_info
65
 
66
  def predict_cap(text, language):