Add support for local LLMs

.gitignore

@@ -131,3 +131,7 @@ dmypy.json
 local/
 *ipynb
 query/
+
+
+*.bin
+*.gguf

README.md

@@ -57,6 +57,26 @@ chem_model = ChemCrow(model="gpt-4-0613", temp=0.1, streaming=False)
 chem_model.run("What is the molecular weight of tylenol?")
 ```
 
+### 💻 Running using local LLMs.
+
+ChemCrow also supports the use of local LLMs, powered by GPT4All, which can be run on a laptop.
+
+A list of supported models is provided [here](https://gpt4all.io/index.html).
+
+```python
+from chemcrow.agents import ChemCrow
+
+chem_model = ChemCrow(
+    model="./models/mistral-7b-instruct-v0.1.Q4_0.gguf",
+    tools_model="./models/mistral-7b-instruct-v0.1.Q4_0.gguf",
+    temp=0.1, verbose=False, max_tokens=100, n_ctx=2048
+)
+output = chem_model.run("What is the molecular weight of tylenol?")
+
+>>> output
+>>> The molecular weight of acetaminophen is 151.17 g/mol ...
+```
+
 ## ✅ Citation
 Bran, Andres M., et al. "ChemCrow: Augmenting large-language models with chemistry tools." arXiv preprint arXiv:2304.05376 (2023).
 

chemcrow/agents/chemcrow.py

@@ -1,39 +1,46 @@
-from typing import Optional
-
-import langchain
+import os
 from dotenv import load_dotenv
+from typing import Optional, Dict
+import langchain
+import nest_asyncio
 from langchain import PromptTemplate, chains
 from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
 from pydantic import ValidationError
 from rmrkl import ChatZeroShotAgent, RetryAgentExecutor
 
+from langchain.llms import GPT4All
+
 from .prompts import FORMAT_INSTRUCTIONS, QUESTION_PROMPT, REPHRASE_TEMPLATE, SUFFIX
 from .tools import make_tools
 
 
-def _make_llm(model, temp, api_key, streaming: bool = False):
+def _make_llm(model, temp, verbose, api_key, max_tokens=1000, n_ctx=2048):
     if model.startswith("gpt-3.5-turbo") or model.startswith("gpt-4"):
-        llm = langchain.chat_models.ChatOpenAI(
-            temperature=temp,
-            model_name=model,
-            request_timeout=1000,
-            streaming=streaming,
-            callbacks=[StreamingStdOutCallbackHandler()],
-            openai_api_key=api_key,
-        )
-    elif model.startswith("text-"):
-        llm = langchain.OpenAI(
-            temperature=temp,
-            model_name=model,
-            streaming=streaming,
-            callbacks=[StreamingStdOutCallbackHandler()],
-            openai_api_key=api_key,
-        )
+        load_dotenv()
+        try:
+            llm = langchain.chat_models.ChatOpenAI(
+                temperature=temp,
+                model_name=model,
+                request_timeout=1000,
+                streaming=True if verbose else False,
+                callbacks=[StreamingStdOutCallbackHandler()] if verbose else [None],
+                openai_api_key = api_key
+            )
+        except:
+            raise ValueError("Invalid OpenAI API key")
+    elif os.path.exists(model):
+        ext = os.path.splitext(model)[-1].lower()
+        if ext == ".gguf":
+            # If GPT4All style weights
+            llm = GPT4All(model=model, max_tokens=max_tokens, verbose=False)
+        else:
+            raise ValueError(f"Found file: {model}, however only models with .gguf format are suported currently.")
     else:
         raise ValueError(f"Invalid model name: {model}")
     return llm
 
 
+
 class ChemCrow:
     def __init__(
         self,
@@ -44,21 +51,23 @@ def __init__(
         max_iterations=40,
         verbose=True,
         streaming: bool = True,
-        openai_api_key: Optional[str] = None,
-        api_keys: dict = {},
+        openai_api_key: str = '',
+        api_keys: Dict[str, str] = {},
+        max_tokens: int = 1000, # Not required for using OpenAI's API
+        n_ctx: int =  2048
     ):
         """Initialize ChemCrow agent."""
 
-        load_dotenv()
-        try:
-            self.llm = _make_llm(model, temp, openai_api_key, streaming)
-        except ValidationError:
-            raise ValueError("Invalid OpenAI API key")
+        self.llm = _make_llm(model, temp, verbose, openai_api_key, max_tokens, n_ctx)
 
         if tools is None:
             api_keys["OPENAI_API_KEY"] = openai_api_key
-            tools_llm = _make_llm(tools_model, temp, openai_api_key, streaming)
-            tools = make_tools(tools_llm, api_keys=api_keys, verbose=verbose)
+            tools_llm = _make_llm(tools_model, temp, verbose, openai_api_key, max_tokens, n_ctx)
+            tools = make_tools(
+                tools_llm,
+                api_keys = api_keys,
+                verbose=verbose
+            )
 
         # Initialize agent
         self.agent_executor = RetryAgentExecutor.from_agent_and_tools(