Update kpe_ranker.py

kpe_ranker.py

@@ -10,6 +10,9 @@ class KpeRanker:
         model_name = os.environ.get("MODEL_NAME")
         model_repo = os.environ.get("MODEL_REPO")
         model_token = os.environ.get("MODEL_TOKEN")
+        ner_model = os.environ.get("NER_MODEL")
+        transformer_model = os.environ.get("TRANSFORMER_MODEL")
+        
         
         local_dir = "./"
         model_path = os.path.join(local_dir, model_name)
@@ -17,8 +20,8 @@ class KpeRanker:
             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')
-        self.ranker_transformer = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2',  device='cpu')
+        self.kpe = KPE(trained_kpe_model= TRAINED_MODEL_ADDR, flair_ner_model= ner_model , device='cpu')
+        self.ranker_transformer = SentenceTransformer(transformer_model,  device='cpu')
 
  
     def extract(self, text, count, using_ner, return_sorted):