Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,6 +2,10 @@ import streamlit as st
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import json
+import os
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+MODEL_DIR = os.path.join(BASE_DIR, "arxiv_dir")
 
 st.set_page_config(
     page_title="Arxiv Classifier",
@@ -14,15 +18,15 @@ st.title("Arxiv Classifier")
 
 @st.cache_resource
 def load_model():
-    model = AutoModelForSequenceClassification.from_pretrained("./arxiv_dir")
-    tokenizer = AutoTokenizer.from_pretrained("./arxiv_dir")
-    with open("./arxiv_dir/id2tag.json") as f:
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+    with open(os.path.join(MODEL_DIR, "id2tag.json")) as f:
         id2tag = {int(k): v for k, v in json.load(f).items()}
-    with open("./arxiv_dir/tag2name.json") as f:
+    with open(os.path.join(MODEL_DIR, "tag2name.json")) as f:
         tag2name = json.load(f)
     model.eval()
     return model, tokenizer, id2tag, tag2name
-
+    
 model, tokenizer, id2tag, tag2name = load_model()
 
 def predict_top95(title, summary=None):