Update kpe_ranker.py

kpe_ranker.py

@@ -7,10 +7,14 @@ from huggingface_hub import hf_hub_download
 
 class KpeRanker:
     def __init__(self):
+        model_name = os.environ.get("MODEL_NAME")
+        model_repo = os.environ.get("MODEL_REPO")
+        model_token = os.environ.get("MODEL_TOKEN")
+        
         local_dir = "./"
-        model_path = os.path.join(local_dir, 'trained_model_10000.pt')
+        model_path = os.path.join(local_dir, model_name)
         if not os.path.isfile(model_path):
-            hf_hub_download(repo_id="ahdsoft/persian-keyphrase-extraction-model", filename="trained_model_10000.pt", local_dir=local_dir, token="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")
+            hf_hub_download(repo_id=model_repo, filename=model_name, local_dir=local_dir, token=model_token)
         TRAINED_MODEL_ADDR = model_path
         # TRAINED_MODEL_ADDR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'trained_model', 'trained_model_10000.pt')
         self.kpe = KPE(trained_kpe_model= TRAINED_MODEL_ADDR, flair_ner_model='flair/ner-english-ontonotes-large', device='cpu')