🐛 Add model-eval to smxm and gliner

Signed-off-by: Marcos Martinez <[email protected]>
IBM · Nov 24, 2024 · 2268aaa · 2268aaa
1 parent f5a2a9a
commit 2268aaa
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 1 deletion.
diff --git a/zshot/evaluation/zshot_evaluate.py b/zshot/evaluation/zshot_evaluate.py
@@ -78,8 +78,10 @@ def make_table(data: Dict, title: str = ""):
         return table
 
     tables = []
-    mode = evaluation.pop('evaluation_mode')
+    mode = evaluation.get('evaluation_mode')
     for component in evaluation:
+        if component == 'evaluation_mode':
+            continue
         # General evaluation
         t_repr = make_table(evaluation[component], f"{component} - {name} \n General - {mode}-based").get_string()
         tables.append(fix_table_title(t_repr))

diff --git a/zshot/linker/linker_gliner.py b/zshot/linker/linker_gliner.py
@@ -35,6 +35,7 @@ def load_models(self):
         """ Load GLINER model """
         if self.model is None:
             self.model = GLiNER.from_pretrained(self.model_name, cache_dir=MODELS_CACHE_PATH).to(self.device)
+            self.model.eval()
 
     def predict(self, docs: Iterator[Doc], batch_size: Optional[Union[int, None]] = None) -> List[List[Span]]:
         """

diff --git a/zshot/linker/linker_smxm.py b/zshot/linker/linker_smxm.py
@@ -39,6 +39,7 @@ def load_models(self):
             self.model = BertTaggerMultiClass.from_pretrained(
                 self.model_name, output_hidden_states=True, cache_dir=MODELS_CACHE_PATH
             ).to(self.device)
+            self.model.eval()
 
     def predict(self, docs: Iterator[Doc], batch_size: Optional[Union[int, None]] = None) -> List[List[Span]]:
         """

diff --git a/zshot/mentions_extractor/mentions_extractor_gliner.py b/zshot/mentions_extractor/mentions_extractor_gliner.py
@@ -30,6 +30,7 @@ def load_models(self):
         """ Load GLINER model """
         if self.model is None:
             self.model = GLiNER.from_pretrained(self.model_name, cache_dir=MODELS_CACHE_PATH).to(self.device)
+            self.model.eval()
 
     def predict(self, docs: Iterator[Doc], batch_size: Optional[Union[int, None]] = None) -> List[List[Span]]:
         """

diff --git a/zshot/mentions_extractor/mentions_extractor_smxm.py b/zshot/mentions_extractor/mentions_extractor_smxm.py
@@ -34,6 +34,7 @@ def load_models(self):
             self.model = BertTaggerMultiClass.from_pretrained(
                 self.model_name, output_hidden_states=True, cache_dir=MODELS_CACHE_PATH
             ).to(self.device)
+            self.model.eval()
 
     def predict(self, docs: Iterator[Doc], batch_size: Optional[Union[int, None]] = None) -> List[List[Span]]:
         """