Fix issues On relik import, OverflowError: Python int too large to co…

…nvert to C long #14 and GoldenRetriever example from README does not work #12

relik/retriever/pytorch_modules/model.py

@@ -265,7 +265,7 @@ def retrieve(
         input_ids: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         token_type_ids: Optional[torch.Tensor] = None,
-        k: int | None = None,
+        k: int = 100,
         max_length: int | None = None,
         precision: Optional[Union[str, int]] = None,
         collate_fn: Optional[Callable] = None,
@@ -306,6 +306,10 @@ def retrieve(
         Returns:
             `List[List[RetrievedSample]]`: The retrieved passages and their indices.
         """
+
+        if "top_k" in kwargs:
+            k = kwargs.pop("top_k")
+
         if self.document_index is None:
             raise ValueError(
                 "The indexer must be indexed before it can be used within the retriever."
@@ -378,6 +382,8 @@ def default_collate_fn(
     ) -> ModelInputs:
         # get text and text pair
         # TODO: check if only retriever is used
+        if not isinstance(x, list):
+            x = [x]
         _text = [sample[0] for sample in x]
         _text_pair = [sample[1] for sample in x]
         _text_pair = None if any([t is None for t in _text_pair]) else _text_pair

setup.py

@@ -130,7 +130,7 @@ def fix_url_dependencies(req: str) -> str:
     ],
     install_requires=install_requirements,
     extras_require=extras,
-    python_requires=">=3.10",
+    python_requires=">=3.10,<3.12",
     find_links=find_links,
     entry_points={
         "console_scripts": ["relik = relik.cli.cli:app"],