finding bottlenecks pt1

interfaces/manifesto.py

@@ -1,3 +1,4 @@
+import time
 import gradio as gr
 
 import os
@@ -9,6 +10,18 @@ from huggingface_hub import HfApi
 
 from label_dicts import MANIFESTO_LABEL_NAMES
 
+class RuntimeMeasure:
+    def __enter__(self):
+        self.start_time = time.time()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        end_time = time.time()
+        runtime = end_time - self.start_time
+        gr.Info(f"Runtime: {runtime} seconds")
+def m(msg):
+    return RuntimeMeasure(msg)
+
 HF_TOKEN = os.environ["hf_read"]
 
 languages = [
@@ -24,23 +37,30 @@ def build_huggingface_path(language: str):
 
 def predict(text, model_id, tokenizer_id):
     device = torch.device("cpu")
-    model = AutoModelForSequenceClassification.from_pretrained(model_id, token=HF_TOKEN)
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
-    model.to(device)
-
-    inputs = tokenizer(text,
-                       max_length=512,
-                       truncation=True,
-                       padding="do_not_pad",
-                       return_tensors="pt").to(device)
-    model.eval()
-
-    with torch.no_grad():
-        logits = model(**inputs).logits
-
-    probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
-    output_pred = {f"[{model.config.id2label[i]}] {MANIFESTO_LABEL_NAMES[int(model.config.id2label[i])]}": probs[i] for i in np.argsort(probs)[::-1]}
-    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>'
+    with m("Loading model"):
+        model = AutoModelForSequenceClassification.from_pretrained(model_id, low_cpu_mem_usage=True, device_map="auto", token=HF_TOKEN)
+    with m("Loading tokenizer"):
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+
+    with m("Tokenizing"):
+        inputs = tokenizer(text,
+                           max_length=512,
+                           truncation=True,
+                           padding="do_not_pad",
+                           return_tensors="pt").to(device)
+    with m("model.eval()"):
+        model.eval()
+
+    with m("Inference"):
+        with torch.no_grad():
+            logits = model(**inputs).logits
+
+    with m("Softmax"):
+        probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
+
+    with m("Output formatting"):
+        output_pred = {f"[{model.config.id2label[i]}] {MANIFESTO_LABEL_NAMES[int(model.config.id2label[i])]}": probs[i] for i in np.argsort(probs)[::-1]}
+        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>'
     return output_pred, output_info
 
 def predict_cap(text, language):