Update

app.py

@@ -1,3 +1,5 @@
+import requests
+import os
 import streamlit as st
 import pandas as pd
 import numpy as np
@@ -9,17 +11,23 @@ import torch
 from transformers import AutoTokenizer
 import joblib
 from model import MultiLabelDeberta
+from huggingface_hub import hf_hub_download
 
 # ========== Loading model and data ==========
 st.set_page_config(page_title="Tag Predictor", layout="wide")
 
+REPO_ID = "Framby/deberta_multilabel"
+deberta_path = hf_hub_download(
+    repo_id=REPO_ID, filename="deberta_multilabel.pt")
+mlb_path = hf_hub_download(repo_id=REPO_ID, filename="mlb.pkl")
+
 
 @st.cache_resource
 def load_model_and_tokenizer():
-    mlb = joblib.load("mlb.pkl")
+    mlb = joblib.load(mlb_path)
     model = MultiLabelDeberta(num_labels=len(mlb.classes_))
     model.load_state_dict(torch.load(
-        "deberta_multilabel.pt", map_location="cpu"))
+        deberta_path, map_location="cpu", weights_only=False))
     model.eval()
     tokenizer = AutoTokenizer.from_pretrained(
         "microsoft/deberta-v3-base", use_fast=False)
@@ -28,33 +36,6 @@ def load_model_and_tokenizer():
 
 model, tokenizer, mlb = load_model_and_tokenizer()
 
-import os
-import requests
-
-def download_from_gdrive(file_id, dest_path):
-    URL = "https://drive.google.com/uc?export=download"
-    session = requests.Session()
-    response = session.get(URL, params={'id': file_id}, stream=True)
-    token = None
-    for key, value in response.cookies.items():
-        if key.startswith('download_warning'):
-            token = value
-    if token:
-        params = {'id': file_id, 'confirm': token}
-        response = session.get(URL, params=params, stream=True)
-    with open(dest_path, "wb") as f:
-        for chunk in response.iter_content(32768):
-            if chunk:
-                f.write(chunk)
-
-
-if not os.path.exists("deberta_multilabel.pt"):
-    download_from_gdrive("1XE_nJwFJwdZj2-I4gH6kAfGuOBczlRzf", "deberta_multilabel.pt")
-
-
-if not os.path.exists("mlb.pkl"):
-    download_from_gdrive("1M2_AVSu9VxAR9NJg75x3UHxiw-2laNCh", "mlb.pkl")
-    
 # ========== data loading ==========
 
 

requirements.txt

@@ -2,6 +2,7 @@
 pandas>=1.3.0
 numpy>=1.21.0
 requests>=2.31.0
+huggingface_hub
 
 # === Visualization ===
 matplotlib>=3.5.0