Update app.py

app.py

@@ -1,56 +1,56 @@
 import streamlit as st
-import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+# import torch
+# from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 
-def combine_title_summary(title, summary):
-    return "title: " + title + " summary: " + summary
+# def combine_title_summary(title, summary):
+#     return "title: " + title + " summary: " + summary
 
 
-tag2ind = {
-    "bio": 0,
-    "physics": 1,
-    "math": 2,
-    "cs": 3,
-}
+# tag2ind = {
+#     "bio": 0,
+#     "physics": 1,
+#     "math": 2,
+#     "cs": 3,
+# }
 
 
-@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)