Merge branch 'main' into dev

# Conflicts:
#	ontolearner/base/learner.py

Author: HamedBabaei

ontolearner/base/learner.py

@@ -236,15 +236,21 @@ def load(self, model_id: str) -> None:
         self.tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side='left', token=self.token)
         self.tokenizer.pad_token = self.tokenizer.eos_token
         if self.device == "cpu":
-            device_map = "cpu"
+            # device_map = "cpu"
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                # device_map=device_map,
+                torch_dtype=torch.bfloat16,
+                token=self.token
+            )
         else:
             device_map = "balanced"
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            device_map=device_map,
-            torch_dtype=torch.bfloat16,
-            token=self.token
-        )
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                device_map=device_map,
+                torch_dtype=torch.bfloat16,
+                token=self.token
+            )
         self.label_mapper.fit()
 
     def generate(self, inputs: List[str], max_new_tokens: int = 50) -> List[str]:
@@ -290,7 +296,8 @@ def generate(self, inputs: List[str], max_new_tokens: int = 50) -> List[str]:
 
         # Decode only the generated part
         decoded_outputs = [self.tokenizer.decode(g, skip_special_tokens=True).strip() for g in generated_tokens]
-
+        print(decoded_outputs)
+        print(self.label_mapper.predict(decoded_outputs))
         # Map the decoded text to labels
         return self.label_mapper.predict(decoded_outputs)
 
@@ -301,9 +308,6 @@ class AutoRetriever(ABC):
     This class defines the interface for retrieval components used in ontology learning.
     Retrievers are responsible for finding semantically similar examples from training
     data to provide context for language models or to make direct predictions.
-
-    Attributes:
-        model: The loaded retrieval/embedding model instance.
     """
 
     def __init__(self) -> None:
@@ -313,7 +317,6 @@ def __init__(self) -> None:
         Sets up the basic structure with a model attribute that will be
         populated when load() is called.
         """
-        self.model: Optional[Any] = None
         self.embedding_model = None
         self.documents = []
         self.embeddings = None

ontolearner/learner/label_mapper.py

@@ -85,6 +85,6 @@ def predict(self, X: List[str]) -> List[str]:
         Returns:
             List[str]: Predicted labels.
         """
-        predictions = list(self.model.predict(X))
+        predictions = self.model.predict(X).tolist()
         self.validate_predicts(predictions)
         return predictions

ontolearner/learner/retriever.py

@@ -66,7 +66,16 @@ def _taxonomy_discovery(self, data: Any, test: bool = False) -> Optional[Any]:
             taxonomic_pairs = [{"parent": candidate, "child": query}
                                for query, candidates in zip(data, candidates_lst)
                                for candidate in candidates if candidate.lower() != query.lower()]
-            return taxonomic_pairs
+            taxonomic_pairs += [{"parent": query, "child": candidate}
+                               for query, candidates in zip(data, candidates_lst)
+                               for candidate in candidates if candidate.lower() != query.lower()]
+            unique_taxonomic_pairs, seen = [], set()
+            for pair in taxonomic_pairs:
+                key = (pair["parent"].lower(), pair["child"].lower()) # Directional key (parent, child)
+                if key not in seen:
+                    seen.add(key)
+                    unique_taxonomic_pairs.append(pair)
+            return unique_taxonomic_pairs
         else:
             warnings.warn("No requirement for fiting the taxonomy discovery model, the predict module will use the input data to do the fit as well.")
 
@@ -86,11 +95,23 @@ def _non_taxonomic_re(self, data: Any, test: bool = False) -> Optional[Any]:
             candidates_lst = self._retriever_predict(data=data['types'], top_k=self.top_k + 1)
             taxonomic_pairs = []
             taxonomic_pairs_query = []
+            seen = set()
             for query, candidates in zip(data['types'], candidates_lst):
                 for candidate in candidates:
                     if candidate != query:
-                        taxonomic_pairs.append((query, candidate))
-                        taxonomic_pairs_query.append(f"Head: {query} \n Tail: {candidate}")
+                        # Directional pair 1: query -> candidate
+                        key1 = (query.lower(), candidate.lower())
+                        if key1 not in seen:
+                            seen.add(key1)
+                            taxonomic_pairs.append((query, candidate))
+                            taxonomic_pairs_query.append(f"Head: {query}\nTail: {candidate}")
+                        # Directional pair 2: candidate -> query
+                        key2 = (candidate.lower(), query.lower())
+                        if key2 not in seen:
+                            seen.add(key2)
+                            taxonomic_pairs.append((candidate, query))
+                            taxonomic_pairs_query.append(f"Head: {candidate}\nTail: {query}")
+
             self._retriever_fit(data=data['relations'])
             candidate_relations_lst = self._retriever_predict(data=taxonomic_pairs_query, top_k=self.top_k)
             non_taxonomic_re = [{"head": head, "tail": tail, "relation": relation}