adding full or peft alignment (#48)

* adding full or peft alignment * formatting --------- Co-authored-by: Shahrzad <[email protected]>
arcee-ai · Jun 18, 2024 · 0079369 · 0079369
1 parent 10e01e3
commit 0079369
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/arcee/__init__.py b/arcee/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "1.2.0"
+__version__ = "1.2.1"
 
 import os
 

diff --git a/arcee/api.py b/arcee/api.py
@@ -278,7 +278,9 @@ def delete_corpus(corpus: str) -> Dict[str, str]:
     return make_request("post", Route.pretraining + "/deleteCorpus", data)
 
 
-def start_alignment(alignment_name: str, qa_set: str, pretrained_model: str) -> Dict[str, str]:
+def start_alignment(
+    alignment_name: str, qa_set: str, pretrained_model: str, full_or_peft: Optional[str] = "full"
+) -> Dict[str, str]:
     """
     Start alignment of a model
 
@@ -288,7 +290,12 @@ def start_alignment(alignment_name: str, qa_set: str, pretrained_model: str) ->
         pretrained_model (str): The name of the pretrained model to use
     """
 
-    data = {"alignment_name": alignment_name, "qa_set_name": qa_set, "pretrained_model": pretrained_model}
+    data = {
+        "alignment_name": alignment_name,
+        "qa_set_name": qa_set,
+        "pretrained_model": pretrained_model,
+        "full_or_peft": full_or_peft,
+    }
 
     return make_request("post", Route.alignment + "/startAlignment", data)