Update inference.py

inference.py

@@ -2,10 +2,8 @@ import streamlit as st
 import torch
 from transformers import BertForTokenClassification, BertTokenizerFast  # Import BertTokenizerFast
 
-def load_model(model_name='linkbert.pth'):
-    model_path = model_name
-    model = BertForTokenClassification.from_pretrained('bert-base-uncased', num_labels=2)
-    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+def load_model(model_name='dejanseo/LinkBERT'):
+    model = BertForTokenClassification.from_pretrained(model_name, num_labels=2)
     model.eval()  # Set the model to inference mode
     return model
 
@@ -45,7 +43,7 @@ def predict_and_annotate(model, tokenizer, text):
 st.title("BERT Token Classification for Anchor Text Prediction")
 
 # Load the model and tokenizer
-model = load_model('linkbert.pth')
+model = load_model('dejanseo/LinkBERT')
 tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')  # Use BertTokenizerFast
 
 # User input text area