Update README.md

README.md

@@ -52,8 +52,8 @@ Then load and use the model:
 ```python
 import torch
 import json
-from torchcrf import CRF
 from torch import nn
+from torchcrf import CRF
 from huggingface_hub import hf_hub_download
 
 # Define the model class (must match training)
@@ -78,47 +78,61 @@ class BiLSTMCRF(nn.Module):
         else:
             return self.crf.decode(emissions, mask=mask.bool())
 
-# Load files from Hugging Face
-config_path = hf_hub_download("TatarNLPWorld/lstm-tatar-morph", "config.json")
-word2id_path = hf_hub_download("TatarNLPWorld/lstm-tatar-morph", "word2id.json")
-weights_path = hf_hub_download("TatarNLPWorld/lstm-tatar-morph", "best_model.pt")
-id2tag_path = hf_hub_download("TatarNLPWorld/lstm-tatar-morph", "id2tag.json")
+# Download required files from Hugging Face
+repo_id = "TatarNLPWorld/lstm-tatar-morph"
+config_path = hf_hub_download(repo_id, "config.json")
+word2id_path = hf_hub_download(repo_id, "word2id.json")
+weights_path = hf_hub_download(repo_id, "best_model.pt")
+id2tag_path = hf_hub_download(repo_id, "id2tag.json")
 
 # Load hyperparameters
 with open(config_path) as f:
     config = json.load(f)
+
 with open(word2id_path) as f:
     word2id = json.load(f)
+
 with open(id2tag_path) as f:
     id2tag = {int(k): v for k, v in json.load(f).items()}
 
 # Instantiate model and load weights
-model = BiLSTMCRF(len(word2id), config['embedding_dim'], config['hidden_dim'], config['num_labels'], config['dropout'])
+model = BiLSTMCRF(
+    vocab_size=len(word2id),
+    emb_dim=config['embedding_dim'],
+    hid_dim=config['hidden_dim'],
+    num_tags=config['num_labels'],
+    dropout=config.get('dropout', 0.5)
+)
 model.load_state_dict(torch.load(weights_path, map_location='cpu'), strict=False)
 model.eval()
 
-def predict(tokens):
+def predict(tokens, max_len=128):
     ids = [word2id.get(w, word2id['<UNK>']) for w in tokens]
     mask = [1] * len(ids)
     orig_len = len(ids)
-    if len(ids) < 128:
-        ids += [0] * (128 - len(ids))
-        mask += [0] * (128 - len(mask))
+    
+    if len(ids) > max_len:
+        ids = ids[:max_len]
+        mask = mask[:max_len]
+        tokens = tokens[:max_len]
     else:
-        ids = ids[:128]
-        mask = mask[:128]
-    input_ids = torch.tensor([ids])
-    mask_tensor = torch.tensor([mask])
+        ids += [0] * (max_len - len(ids))
+        mask += [0] * (max_len - len(mask))
+    
+    input_ids = torch.tensor([ids], dtype=torch.long)
+    mask_tensor = torch.tensor([mask], dtype=torch.long)
+    
     with torch.no_grad():
         preds = model(input_ids, mask_tensor)[0]
+    
     preds = preds[:orig_len]
     return [id2tag[p] for p in preds]
 
 # Example
 tokens = ["Татар", "теле", "бик", "бай", "."]
 tags = predict(tokens)
-for t, tag in zip(tokens, tags):
-    print(f"{t} -> {tag}")
+for token, tag in zip(tokens, tags):
+    print(f"{token} -> {tag}")
 ```
 
 Expected output: