Update app.py

app.py

@@ -1,11 +1,12 @@
 import streamlit as st
 import torch
-from safetensors.torch import load_file
 from transformers import AutoConfig, AutoTokenizer, AutoModel
-from ST2ModelV2_6 import ST2ModelV2
 from huggingface_hub import login
 import re
 import copy
+from src.models.modeling_st2 import ST2ModelV2, SignalDetector
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
 hf_token = st.secrets["HUGGINGFACE_TOKEN"]
 login(token=hf_token)
@@ -16,10 +17,9 @@ login(token=hf_token)
 @st.cache_resource
 def load_model():
     
-    model_name = "anamargarida/Final"  
-    
-    config = AutoConfig.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    config = AutoConfig.from_pretrained("roberta-large")
+
+    tokenizer = AutoTokenizer.from_pretrained("roberta-large", use_fast=True, add_prefix_space=True)
     
     class Args:
         def __init__(self):
@@ -30,9 +30,19 @@ def load_model():
 
     args = Args()
 
-    # Load the model directly from Hugging Face
-    model = ST2ModelV2.from_pretrained(model_name, config=config, args=args)
+    model = ST2ModelV2(args)
+
+
+    repo_id = "anamargarida/SpanExtractionWithSignalCls_2"  
+    filename = "model.safetensors"  
+
+    # Download the model file
+    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+
+    # Load the model weights
+    state_dict = load_file(model_path)
 
+    model.load_state_dict(state_dict)
     
     return tokenizer, model