Update interfaces/cap.py

interfaces/cap.py

@@ -95,28 +95,28 @@ def build_huggingface_path(language: str, domain: str):
 
 def predict(text, model_id, tokenizer_id):
     device = torch.device("cpu")
-    gradio.Info("Loading model")
+    gr.Info("Loading model")
     model = AutoModelForSequenceClassification.from_pretrained(model_id, token=HF_TOKEN)
-    gradio.Info("Loading tokenizer")
+    gr.Info("Loading tokenizer")
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
-    gradio.Info("Mapping model to device")
+    gr.Info("Mapping model to device")
     model.to(device)
 
-    gradio.Info("Tokenizing")
+    gr.Info("Tokenizing")
     inputs = tokenizer(text,
                        max_length=4,
                        truncation=True,
                        padding="do_not_pad",
                        return_tensors="pt").to(device)
 
-    gradio.Info("model.eval()")
+    gr.Info("model.eval()")
     model.eval()
 
-    gradio.Info("Prediction")
+    gr.Info("Prediction")
     with torch.no_grad():
         logits = model(**inputs).logits
 
-    gradio.Info("Softmax")
+    gr.Info("Softmax")
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
     output_pred = {f"[{CAP_NUM_DICT[i]}] {CAP_LABEL_NAMES[CAP_NUM_DICT[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>'