Merge pull request #43 from arcee-ai/add-lora-dropout-to-retriever

Add lora dropout and remove non-required input param
arcee-ai · Sep 20, 2023 · 7614a21 · 7614a21
2 parents 0cf02bc + c94c68a
commit 7614a21
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/dalm/models/retriever_only_base_model.py b/dalm/models/retriever_only_base_model.py
@@ -72,13 +72,15 @@ def __get_bnb_config() -> BitsAndBytesConfig:
     def __get_lora_config(
         r: int = 8,
         lora_alpha: int = 16,
+        lora_dropout: float = 0.05,
         bias: str = "none",
         target_modules: Optional[Union[List[str], str]] = None,
     ) -> LoraConfig:
         return LoraConfig(
             task_type=TaskType.FEATURE_EXTRACTION,
             r=r,
             lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
             bias=bias,
             target_modules=target_modules,
         )
diff --git a/dalm/training/retriever_only/train_retriever_only.py b/dalm/training/retriever_only/train_retriever_only.py
@@ -211,7 +211,6 @@ def main() -> None:
             tokenizer,
             query_col_name=args.dataset_query_col_name,
             passage_col_name=args.dataset_passage_col_name,
-            max_length=args.max_length,
             query_max_len=args.query_max_len,
             passage_max_len=args.passage_max_len,
         ),