Update app.py

app.py

@@ -15,42 +15,42 @@ tag2ind = {
 }
 
 
-# @st.cache_resource
-# def load_model():
-#     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+@st.cache_resource
+def load_model():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-#     # assert torch.cuda.is_available()
-#     tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-cased")
-#     model = AutoModelForSequenceClassification.from_pretrained(
-#         "./my_model/checkpoint-513"
-#     ).to(device)
+    # assert torch.cuda.is_available()
+    tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-cased")
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "./my_model/checkpoint-513"
+    ).to(device)
 
-#     return tokenizer, model
+    return tokenizer, model
 
 
-# tokenizer, model = load_model()
+tokenizer, model = load_model()
 
 
-# def run_model(model, tokenizer, title, summary):
-#     text = combine_title_summary(title, summary)
+def run_model(model, tokenizer, title, summary):
+    text = combine_title_summary(title, summary)
 
-#     tokens_info = tokenizer(
-#         text,
-#         padding=False,
-#         truncation=True,
-#         return_tensors="pt",
-#     )
+    tokens_info = tokenizer(
+        text,
+        padding=False,
+        truncation=True,
+        return_tensors="pt",
+    )
 
-#     model.eval()
-#     model.cpu()
-#     with torch.no_grad():
-#         out = model(**tokens_info)
-#         probs = torch.nn.functional.softmax(out.logits, dim=-1)[0]
+    model.eval()
+    model.cpu()
+    with torch.no_grad():
+        out = model(**tokens_info)
+        probs = torch.nn.functional.softmax(out.logits, dim=-1)[0]
 
-#         result = f"Text: `{text}`\nPrediction (prob): \n" + "\n".join(
-#             [f"{tag}={tag_prob}" for tag, tag_prob in zip(tag2ind, probs)]
-#         )
-#         return result
+        result = f"Text: `{text}`\nPrediction (prob): \n" + "\n".join(
+            [f"{tag}={tag_prob}" for tag, tag_prob in zip(tag2ind, probs)]
+        )
+        return result
 
 
 title = st.text_input(label="Title", value="")
@@ -62,5 +62,5 @@ if st.button("Submit"):
         result = combine_title_summary(title, abstract)
         st.success(result)
 
-        # result = run_model(model, tokenizer, title, abstract)
-        # st.success(result)
+        result = run_model(model, tokenizer, title, abstract)
+        st.success(result)