Update model.py

model.py

@@ -5,10 +5,14 @@ import sys
 
 from tokenizers import Tokenizer
 
-MODEL_PATH = './model.pth'
-TOKENIZER_PATH = './hindi-english_bpe_tokenizer.json'
+from huggingface_hub import hf_hub_download
+from tokenizers import Tokenizer
+
+model_file = hf_hub_download(repo_id="Kush26/Transformer_Translation", filename="model.pth")
+tokenizer_file = hf_hub_download(repo_id="Kush26/Transformer_Translation", filename="hindi-english_bpe_tokenizer.json")
+
+tokenizer = Tokenizer.from_file(tokenizer_file)
 
-tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
 vocab_size = tokenizer.get_vocab_size()
 pad_token_id = tokenizer.token_to_id('[PAD]')
 
@@ -313,10 +317,7 @@ model = BuildTransformer(vocab_size,
                          config["dropout"],
                          config["d_ff"]).to(device)
 
-# total_parameters = sum(p.numel() for p in model.parameters())
-# print(f"Totoal Parameters = {total_parameters}")
-
-checkpoint = torch.load(MODEL_PATH, map_location=device)
+checkpoint = torch.load(model_file, map_location=device)
 model.load_state_dict(checkpoint['model_state_dict'])
 model.eval()