add local mode and readme updates

README.md

@@ -68,7 +68,14 @@ df['capitol'] = pd.DataFrame({'State': ['Colorado', 'Kansas', 'California', 'New
 
 ## Sketch currently uses `prompts.approx.dev` to help run with minimal setup
 
-In the future, we plan to update the prompts at this endpoint with our own custom foundation model, built to answer questions more accurately than GPT-3 can with its minimal data context. 
+You can also directly use a few pre-built hugging face models (right now `MPT-7B` and `StarCoder`), which will run entirely locally (once you download the model weights from HF).
+Do this by setting environment 3 variables:
+
+```python
+os.environ['LAMBDAPROMPT_BACKEND'] = 'StarCoder'
+os.environ['SKETCH_USE_REMOTE_LAMBDAPROMPT'] = 'False'
+os.environ['HF_ACCESS_TOKEN'] = 'your_hugging_face_token'
+```
 
 You can also directly call OpenAI directly (and not use our endpoint) by using your own API key. To do this, set 2 environment variables.
 

pyproject.toml

@@ -23,6 +23,9 @@ dependencies = [
 urls = {homepage = "https://github.com/approximatelabs/sketch"}
 dynamic = ["version"]
 
+[project.optional-dependencies]
+local = ["lambdaprompt[local]"]
+all = ["sketch[local]"]
 
 [tool.setuptools_scm]
 

sketch/pandas_extension.py

@@ -181,7 +181,7 @@ def call_prompt_on_dataframe(df, prompt, **kwargs):
  return text_to_copy
 
 
-howto_prompt = lambdaprompt.GPT3Prompt(
+howto_prompt = lambdaprompt.Completion(
  """
 For the pandas dataframe ({{ dfname }}) the user wants code to solve a problem.
 Summary statistics and descriptive data of dataframe [`{{ dfname }}`]:
@@ -234,7 +234,7 @@ def howto_from_parts(
  return code
 
 
-ask_prompt = lambdaprompt.GPT3Prompt(
+ask_prompt = lambdaprompt.Completion(
  """
 For the pandas dataframe ({{ dfname }}) the user wants an answer to a question about the data.
 Summary statistics and descriptive data of dataframe [`{{ dfname }}`]:
@@ -338,7 +338,7 @@ def apply(self, prompt_template_string, **kwargs):
  raise RuntimeError(
  f"Too many rows for apply \n (SKETCH_ROW_OVERRIDE_LIMIT: {row_limit}, Actual: {len(self._obj)})"
  )
- new_gpt3_prompt = lambdaprompt.GPT3Prompt(prompt_template_string)
+ new_gpt3_prompt = lambdaprompt.Completion(prompt_template_string)
  named_args = new_gpt3_prompt.get_named_args()
  known_args = set(self._obj.columns) | set(kwargs.keys())
  needed_args = set(named_args)