Add BiLSTMConfig and update BiLSTMAttentionBERT for improved model configuration and loading

utils/model.py

@@ -3,6 +3,8 @@ import torch.nn as nn
 from transformers import PreTrainedModel, AutoModel, PretrainedConfig
 
 class BiLSTMConfig(PretrainedConfig):
+    model_type = "bilstm_attention"
+    
     def __init__(self, hidden_dim=128, num_classes=22, num_layers=2, dropout=0.5, **kwargs):
         super().__init__(**kwargs)
         self.hidden_dim = hidden_dim
@@ -11,12 +13,20 @@ class BiLSTMConfig(PretrainedConfig):
         self.dropout = dropout
 
 class BiLSTMAttentionBERT(PreTrainedModel):
+    config_class = BiLSTMConfig
+    base_model_prefix = "bilstm_attention"
+    
     def __init__(self, config):
         super().__init__(config)
         self.config = config
         self.bert = AutoModel.from_pretrained('dmis-lab/biobert-base-cased-v1.2')
-        self.lstm = nn.LSTM(768, config.hidden_dim, config.num_layers, 
-                           batch_first=True, bidirectional=True)
+        self.lstm = nn.LSTM(
+            768, 
+            config.hidden_dim,
+            config.num_layers,
+            batch_first=True,
+            bidirectional=True
+        )
         self.dropout = nn.Dropout(config.dropout)
         self.fc = nn.Linear(config.hidden_dim * 2, config.num_classes)
         

utils/prediction.py

@@ -1,16 +1,31 @@
-from utils.model import BiLSTMAttentionBERT
+from utils.model import BiLSTMAttentionBERT, BiLSTMConfig
 import torch
 from transformers import AutoTokenizer
 from sklearn.preprocessing import LabelEncoder
 import numpy as np
 import streamlit as st
 import requests
+from huggingface_hub import hf_hub_download
 
 
 
 def load_model_for_prediction():
     try:
         st.write("Starting model loading...")
+        config = BiLSTMConfig(
+            hidden_dim=128,
+            num_classes=22,
+            num_layers=2,
+            dropout=0.5
+        )
+        
+        model = BiLSTMAttentionBERT(config)
+        model_path = hf_hub_download(
+            repo_id="joko333/BiLSTM_v01",
+            filename="model_epoch8_acc72.53.pt"
+        )
+        state_dict = torch.load(model_path, map_location='cpu')
+        model.load_state_dict(state_dict)
         
         # Test Hugging Face connectivity
         st.write("Testing connection to Hugging Face...")