Merge pull request #17 from wangyuxinwhy/fix-infer-record-type

🐞fix:infer record type bug

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "uniem"
-version = "0.2.2"
+version = "0.2.3"
 description = "unified embedding model"
 authors = ["wangyuxin <[email protected]>"]
 license = "MIT"

tests/test_data_structures.py

@@ -0,0 +1,31 @@
+import pytest
+from uniem.data_structures import RecordType, infer_record_type
+
+
+@pytest.mark.parametrize(
+    'record, expected_record_type',
+    [
+        pytest.param(
+            {'text': 'I like apples', 'text_pos': 'I like oranges'},
+            RecordType.PAIR,
+            id='pair',
+        ),
+        pytest.param(
+            {'text': 'I like apples', 'text_pos': 'I like oranges', 'source': 'wikipedia'},
+            RecordType.PAIR,
+            id='pair_with_extra_fields',
+        ),
+        pytest.param(
+            {'text': 'I like apples', 'text_pos': 'I like oranges', 'text_neg': 'I want to eat apples'},
+            RecordType.TRIPLET,
+            id='triplet',
+        ),
+        pytest.param(
+            {'sentence1': 'I like apples', 'sentence2': 'I like oranges', 'label': 1.0},
+            RecordType.SCORED_PAIR,
+            id='scored_pair',
+        ),
+    ],
+)
+def test_infer_record_type(record: dict, expected_record_type: RecordType):
+    assert infer_record_type(record) == expected_record_type

uniem/data.py

@@ -13,7 +13,7 @@
     RecordType,
     ScoredPairRecord,
     TripletRecord,
-    get_record_type,
+    infer_record_type,
     record_type_cls_map,
 )
 from uniem.types import Tokenizer
@@ -140,7 +140,7 @@ def __init__(
         if record_type:
             self.record_type = RecordType(record_type)
         else:
-            self.record_type = get_record_type(dataset[0])
+            self.record_type = infer_record_type(dataset[0])
         self.record_cls = record_type_cls_map[self.record_type]
 
     def __getitem__(self, index: int):

uniem/data_structures.py

@@ -29,14 +29,15 @@ class ScoredPairRecord:
     label: float
 
 
+# * Order matters
 record_type_cls_map: dict[RecordType, Any] = {
-    RecordType.PAIR: PairRecord,
-    RecordType.TRIPLET: TripletRecord,
     RecordType.SCORED_PAIR: ScoredPairRecord,
+    RecordType.TRIPLET: TripletRecord,
+    RecordType.PAIR: PairRecord,
 }
 
 
-def get_record_type(record: dict) -> RecordType:
+def infer_record_type(record: dict) -> RecordType:
     record_type_field_names_map = {
         record_type: [field.name for field in fields(record_cls)] for record_type, record_cls in record_type_cls_map.items()
     }

uniem/finetuner.py

@@ -19,7 +19,7 @@
     ScoredPairCollator,
     TripletCollator,
 )
-from uniem.data_structures import RecordType, get_record_type
+from uniem.data_structures import RecordType, infer_record_type
 from uniem.model import (
     EmbedderForPairInBatchNegTrain,
     EmbedderForScoredPairTrain,
@@ -41,6 +41,7 @@ def __init__(
         self,
         model_name_or_path: str,
         dataset: RawDataset,
+        record_type: RecordType | str | None = None,
     ):
         self.model_name_or_path = model_name_or_path
         self.raw_dataset = dataset
@@ -54,7 +55,10 @@ def __init__(
                 self.raw_dataset,
                 None,
             )
-        self.record_type = get_record_type(self.raw_train_dataset[0])  # type: ignore
+
+        record_type = RecordType(record_type) if isinstance(record_type, str) else record_type
+        self.record_type = record_type or infer_record_type(self.raw_train_dataset[0])
+
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)
 
     def create_finetune_datasets(

uniem/version.py

@@ -1 +1 @@
-__version__ = '0.1.1'
+__version__ = '0.2.3'