From 5d513550b7f79dd2e6aec2b9a692c4cbfb5fbc64 Mon Sep 17 00:00:00 2001
From: Saisri Vishwanath <saisrivishwanath@Saisris-MacBook-Pro.local>
Date: Wed, 6 Dec 2023 22:35:54 -0500
Subject: [PATCH] quickstartguided

---
 examples/quickstart_guided.py |  9 +++------
 src/llm_vm/client.py          | 18 ++++++++++++++++--
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/examples/quickstart_guided.py b/examples/quickstart_guided.py
index ed20cac3..0fcd13b0 100644
--- a/examples/quickstart_guided.py
+++ b/examples/quickstart_guided.py
@@ -8,13 +8,10 @@
 client = Client(big_model='pythia', small_model='pythia')
 
 # Put in your prompt and go!
-# response = client.complete(prompt = 'Is 1+1=10000?',
-#                            context='', openai_key=settings.openai_api_key, finetune = False, data_synthesis = False, regex=r"\s*([Yy]es|[Nn]o|[Nn]ever|[Aa]lways)")
+#response = client.complete(prompt = 'Is 1+1=10000?', context='', openai_key=settings.openai_api_key, finetune = False, data_synthesis = False, regex=r"\s*([Yy]es|[Nn]o|[Nn]ever|[Aa]lways)")
 # print(response)
-# response = client.complete(prompt = 'Did MJ win 6 titles with the Bulls',
-#                            context='', openai_key=settings.openai_api_key, finetune = False, data_synthesis = False,choices=["Hell yeah dude that's correct","No way, that's hella false"])
+#response = client.complete(prompt = 'Did MJ win 6 titles with the Bulls', context='', openai_key=settings.openai_api_key, finetune = False, data_synthesis = False,choices=["Hell yeah dude that's correct","No way, that's hella false"])
 # print(response)
-response = client.complete(prompt = 'What does 1+1 equal?',
-                           context='',  openai_key=settings.openai_api_key, finetune = True, data_synthesis = True, type="integer")
+response = client.complete(prompt = 'What does 1+1 equal?',context='',  openai_key=settings.openai_api_key, finetune = True, data_synthesis = True, type="integer")
 print(response)
 
diff --git a/src/llm_vm/client.py b/src/llm_vm/client.py
index 7d605723..798f3aaa 100644
--- a/src/llm_vm/client.py
+++ b/src/llm_vm/client.py
@@ -110,7 +110,10 @@ def complete(self, prompt,
                  stoptoken = None,
                  tools = None,
                  openai_kwargs = {},
-                 hf_kwargs = {}):
+                 hf_kwargs = {},
+                 type = None,
+                 choices = None,
+                 regex = None):
         """
         This function is Anarchy's completion entry point
 
@@ -143,12 +146,23 @@ def complete(self, prompt,
                 kwargs.update({"temperature":temperature})
             if stoptoken is not None:
                 kwargs.update({"stop":stoptoken})
-                
+            if regex is not None:
+                kwargs.update({"regex":regex})
+            if type is not None:
+                kwargs.update({"type":type})
+            if choices is not None:
+                kwargs.update({"choices":choices})
         else:
             kwargs = hf_kwargs
             if temperature > 0:
                 kwargs.update({"do_sample": True})
                 kwargs.update({"temperature":temperature})
+            if regex is not None:
+                kwargs.update({"regex":regex})
+            if type is not None:
+                kwargs.update({"type":type})
+            if choices is not None:
+                kwargs.update({"choices":choices})
 
         if openai_key is not None:
             self.openai_key = openai_key