From 9bb7195085a843db3f53e4fd9a51c79f3698a19d Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Tue, 13 Dec 2022 06:46:01 -0800 Subject: [PATCH] Harrison/llm saving (#331) Co-authored-by: Akash Samant <70665700+asamant21@users.noreply.github.com> --- langchain/llms/__init__.py | 14 ++++- langchain/llms/ai21.py | 7 ++- langchain/llms/base.py | 57 ++++++++++++++++++- langchain/llms/cohere.py | 5 ++ langchain/llms/huggingface_hub.py | 10 +++- langchain/llms/loading.py | 42 ++++++++++++++ langchain/llms/manifest.py | 5 ++ langchain/llms/nlpcloud.py | 5 ++ langchain/llms/openai.py | 7 ++- tests/integration_tests/llms/test_ai21.py | 11 ++++ tests/integration_tests/llms/test_cohere.py | 12 ++++ .../llms/test_huggingface_hub.py | 12 ++++ tests/integration_tests/llms/test_nlpcloud.py | 12 ++++ tests/integration_tests/llms/test_openai.py | 11 ++++ tests/integration_tests/llms/utils.py | 16 ++++++ tests/unit_tests/agents/test_agent.py | 17 ++++-- tests/unit_tests/agents/test_react.py | 23 +++++--- tests/unit_tests/chains/test_natbot.py | 9 ++- tests/unit_tests/llms/fake_llm.py | 17 ++++-- tests/unit_tests/llms/test_loading.py | 15 +++++ 20 files changed, 279 insertions(+), 28 deletions(-) create mode 100644 langchain/llms/loading.py create mode 100644 tests/integration_tests/llms/utils.py create mode 100644 tests/unit_tests/llms/test_loading.py diff --git a/langchain/llms/__init__.py b/langchain/llms/__init__.py index fab84d41d..01e3c702a 100644 --- a/langchain/llms/__init__.py +++ b/langchain/llms/__init__.py @@ -1,7 +1,19 @@ """Wrappers on top of large language models APIs.""" +from typing import Dict, Type + +from langchain.llms.ai21 import AI21 +from langchain.llms.base import LLM from langchain.llms.cohere import Cohere from langchain.llms.huggingface_hub import HuggingFaceHub from langchain.llms.nlpcloud import NLPCloud from langchain.llms.openai import OpenAI -__all__ = ["Cohere", "NLPCloud", "OpenAI", "HuggingFaceHub"] +__all__ = ["Cohere", "NLPCloud", "OpenAI", "HuggingFaceHub", "AI21"] + +type_to_cls_dict: Dict[str, Type[LLM]] = { + "ai21": AI21, + "cohere": Cohere, + "huggingface_hub": HuggingFaceHub, + "nlpcloud": NLPCloud, + "openai": OpenAI, +} diff --git a/langchain/llms/ai21.py b/langchain/llms/ai21.py index a870d9e4b..3678473df 100644 --- a/langchain/llms/ai21.py +++ b/langchain/llms/ai21.py @@ -19,7 +19,7 @@ class AI21PenaltyData(BaseModel): applyToEmojis: bool = True -class AI21(BaseModel, LLM): +class AI21(LLM, BaseModel): """Wrapper around AI21 large language models. To use, you should have the environment variable ``AI21_API_KEY`` @@ -96,6 +96,11 @@ class AI21(BaseModel, LLM): """Get the identifying parameters.""" return {**{"model": self.model}, **self._default_params} + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "ai21" + def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Call out to AI21's complete endpoint. diff --git a/langchain/llms/base.py b/langchain/llms/base.py index 367aa503c..0eb8b73b7 100644 --- a/langchain/llms/base.py +++ b/langchain/llms/base.py @@ -1,6 +1,11 @@ """Base interface for large language models to expose.""" +import json from abc import ABC, abstractmethod -from typing import Any, List, Mapping, NamedTuple, Optional +from pathlib import Path +from typing import Any, Dict, List, Mapping, NamedTuple, Optional, Union + +import yaml +from pydantic import BaseModel, Extra class Generation(NamedTuple): @@ -21,9 +26,14 @@ class LLMResult(NamedTuple): """For arbitrary LLM provider specific output.""" -class LLM(ABC): +class LLM(BaseModel, ABC): """LLM wrapper should take in a prompt and return a string.""" + class Config: + """Configuration for this pydantic object.""" + + extra = Extra.forbid + def generate( self, prompts: List[str], stop: Optional[List[str]] = None ) -> LLMResult: @@ -68,3 +78,46 @@ class LLM(ABC): """Get a string representation of the object for printing.""" cls_name = f"\033[1m{self.__class__.__name__}\033[0m" return f"{cls_name}\nParams: {self._identifying_params}" + + @property + @abstractmethod + def _llm_type(self) -> str: + """Return type of llm.""" + + def _llm_dict(self) -> Dict: + """Return a dictionary of the prompt.""" + starter_dict = dict(self._identifying_params) + starter_dict["_type"] = self._llm_type + return starter_dict + + def save(self, file_path: Union[Path, str]) -> None: + """Save the LLM. + + Args: + file_path: Path to file to save the LLM to. + + Example: + .. code-block:: python + + llm.save(file_path="path/llm.yaml") + """ + # Convert file to Path object. + if isinstance(file_path, str): + save_path = Path(file_path) + else: + save_path = file_path + + directory_path = save_path.parent + directory_path.mkdir(parents=True, exist_ok=True) + + # Fetch dictionary to save + prompt_dict = self._llm_dict() + + if save_path.suffix == ".json": + with open(file_path, "w") as f: + json.dump(prompt_dict, f, indent=4) + elif save_path.suffix == ".yaml": + with open(file_path, "w") as f: + yaml.dump(prompt_dict, f, default_flow_style=False) + else: + raise ValueError(f"{save_path} must be json or yaml") diff --git a/langchain/llms/cohere.py b/langchain/llms/cohere.py index e051ba47f..3853cfb90 100644 --- a/langchain/llms/cohere.py +++ b/langchain/llms/cohere.py @@ -85,6 +85,11 @@ class Cohere(LLM, BaseModel): """Get the identifying parameters.""" return {**{"model": self.model}, **self._default_params} + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "cohere" + def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Call out to Cohere's generate endpoint. diff --git a/langchain/llms/huggingface_hub.py b/langchain/llms/huggingface_hub.py index 5cded6773..74f0545ae 100644 --- a/langchain/llms/huggingface_hub.py +++ b/langchain/llms/huggingface_hub.py @@ -74,7 +74,15 @@ class HuggingFaceHub(LLM, BaseModel): def _identifying_params(self) -> Mapping[str, Any]: """Get the identifying parameters.""" _model_kwargs = self.model_kwargs or {} - return {**{"repo_id": self.repo_id}, **_model_kwargs} + return { + **{"repo_id": self.repo_id, "task": self.task}, + **{"model_kwargs": _model_kwargs}, + } + + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "huggingface_hub" def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Call out to HuggingFace Hub's inference endpoint. diff --git a/langchain/llms/loading.py b/langchain/llms/loading.py new file mode 100644 index 000000000..d5881ba45 --- /dev/null +++ b/langchain/llms/loading.py @@ -0,0 +1,42 @@ +"""Base interface for loading large language models apis.""" +import json +from pathlib import Path +from typing import Union + +import yaml + +from langchain.llms import type_to_cls_dict +from langchain.llms.base import LLM + + +def load_llm_from_config(config: dict) -> LLM: + """Load LLM from Config Dict.""" + if "_type" not in config: + raise ValueError("Must specify an LLM Type in config") + config_type = config.pop("_type") + + if config_type not in type_to_cls_dict: + raise ValueError(f"Loading {config_type} LLM not supported") + + llm_cls = type_to_cls_dict[config_type] + return llm_cls(**config) + + +def load_llm(file: Union[str, Path]) -> LLM: + """Load LLM from file.""" + # Convert file to Path object. + if isinstance(file, str): + file_path = Path(file) + else: + file_path = file + # Load from either json or yaml. + if file_path.suffix == ".json": + with open(file_path) as f: + config = json.load(f) + elif file_path.suffix == ".yaml": + with open(file_path, "r") as f: + config = yaml.safe_load(f) + else: + raise ValueError("File type must be json or yaml") + # Load the LLM from the config now. + return load_llm_from_config(config) diff --git a/langchain/llms/manifest.py b/langchain/llms/manifest.py index 0e1fb7205..665f0c5b4 100644 --- a/langchain/llms/manifest.py +++ b/langchain/llms/manifest.py @@ -37,6 +37,11 @@ class ManifestWrapper(LLM, BaseModel): kwargs = self.llm_kwargs or {} return {**self.client.client.get_model_params(), **kwargs} + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "manifest" + def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Call out to LLM through Manifest.""" if stop is not None and len(stop) != 1: diff --git a/langchain/llms/nlpcloud.py b/langchain/llms/nlpcloud.py index d9e4c54e2..ae3dc07bb 100644 --- a/langchain/llms/nlpcloud.py +++ b/langchain/llms/nlpcloud.py @@ -106,6 +106,11 @@ class NLPCloud(LLM, BaseModel): """Get the identifying parameters.""" return {**{"model_name": self.model_name}, **self._default_params} + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "nlpcloud" + def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Call out to NLPCloud's create endpoint. diff --git a/langchain/llms/openai.py b/langchain/llms/openai.py index fd88d873f..444840685 100644 --- a/langchain/llms/openai.py +++ b/langchain/llms/openai.py @@ -142,7 +142,12 @@ class OpenAI(LLM, BaseModel): @property def _identifying_params(self) -> Mapping[str, Any]: """Get the identifying parameters.""" - return {**{"model": self.model_name}, **self._default_params} + return {**{"model_name": self.model_name}, **self._default_params} + + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "openai" def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Call out to OpenAI's create endpoint. diff --git a/tests/integration_tests/llms/test_ai21.py b/tests/integration_tests/llms/test_ai21.py index 3737ee5c6..16a8716d6 100644 --- a/tests/integration_tests/llms/test_ai21.py +++ b/tests/integration_tests/llms/test_ai21.py @@ -1,6 +1,9 @@ """Test AI21 API wrapper.""" +from pathlib import Path + from langchain.llms.ai21 import AI21 +from langchain.llms.loading import load_llm def test_ai21_call() -> None: @@ -8,3 +11,11 @@ def test_ai21_call() -> None: llm = AI21(maxTokens=10) output = llm("Say foo:") assert isinstance(output, str) + + +def test_saving_loading_llm(tmp_path: Path) -> None: + """Test saving/loading an AI21 LLM.""" + llm = AI21(maxTokens=10) + llm.save(file_path=tmp_path / "ai21.yaml") + loaded_llm = load_llm(tmp_path / "ai21.yaml") + assert llm == loaded_llm diff --git a/tests/integration_tests/llms/test_cohere.py b/tests/integration_tests/llms/test_cohere.py index f1a8a6c37..4c260982b 100644 --- a/tests/integration_tests/llms/test_cohere.py +++ b/tests/integration_tests/llms/test_cohere.py @@ -1,6 +1,10 @@ """Test Cohere API wrapper.""" +from pathlib import Path + from langchain.llms.cohere import Cohere +from langchain.llms.loading import load_llm +from tests.integration_tests.llms.utils import assert_llm_equality def test_cohere_call() -> None: @@ -8,3 +12,11 @@ def test_cohere_call() -> None: llm = Cohere(max_tokens=10) output = llm("Say foo:") assert isinstance(output, str) + + +def test_saving_loading_llm(tmp_path: Path) -> None: + """Test saving/loading an Cohere LLM.""" + llm = Cohere(max_tokens=10) + llm.save(file_path=tmp_path / "cohere.yaml") + loaded_llm = load_llm(tmp_path / "cohere.yaml") + assert_llm_equality(llm, loaded_llm) diff --git a/tests/integration_tests/llms/test_huggingface_hub.py b/tests/integration_tests/llms/test_huggingface_hub.py index aa181a87c..df0b44161 100644 --- a/tests/integration_tests/llms/test_huggingface_hub.py +++ b/tests/integration_tests/llms/test_huggingface_hub.py @@ -1,8 +1,12 @@ """Test HuggingFace API wrapper.""" +from pathlib import Path + import pytest from langchain.llms.huggingface_hub import HuggingFaceHub +from langchain.llms.loading import load_llm +from tests.integration_tests.llms.utils import assert_llm_equality def test_huggingface_text_generation() -> None: @@ -24,3 +28,11 @@ def test_huggingface_call_error() -> None: llm = HuggingFaceHub(model_kwargs={"max_new_tokens": -1}) with pytest.raises(ValueError): llm("Say foo:") + + +def test_saving_loading_llm(tmp_path: Path) -> None: + """Test saving/loading an HuggingFaceHub LLM.""" + llm = HuggingFaceHub(repo_id="gpt2", model_kwargs={"max_new_tokens": 10}) + llm.save(file_path=tmp_path / "hf.yaml") + loaded_llm = load_llm(tmp_path / "hf.yaml") + assert_llm_equality(llm, loaded_llm) diff --git a/tests/integration_tests/llms/test_nlpcloud.py b/tests/integration_tests/llms/test_nlpcloud.py index 9e4664e4f..4c5ccca0d 100644 --- a/tests/integration_tests/llms/test_nlpcloud.py +++ b/tests/integration_tests/llms/test_nlpcloud.py @@ -1,6 +1,10 @@ """Test NLPCloud API wrapper.""" +from pathlib import Path + +from langchain.llms.loading import load_llm from langchain.llms.nlpcloud import NLPCloud +from tests.integration_tests.llms.utils import assert_llm_equality def test_nlpcloud_call() -> None: @@ -8,3 +12,11 @@ def test_nlpcloud_call() -> None: llm = NLPCloud(max_length=10) output = llm("Say foo:") assert isinstance(output, str) + + +def test_saving_loading_llm(tmp_path: Path) -> None: + """Test saving/loading an NLPCloud LLM.""" + llm = NLPCloud(max_length=10) + llm.save(file_path=tmp_path / "nlpcloud.yaml") + loaded_llm = load_llm(tmp_path / "nlpcloud.yaml") + assert_llm_equality(llm, loaded_llm) diff --git a/tests/integration_tests/llms/test_openai.py b/tests/integration_tests/llms/test_openai.py index 4050e42ed..721b5b43d 100644 --- a/tests/integration_tests/llms/test_openai.py +++ b/tests/integration_tests/llms/test_openai.py @@ -1,7 +1,10 @@ """Test OpenAI API wrapper.""" +from pathlib import Path + import pytest +from langchain.llms.loading import load_llm from langchain.llms.openai import OpenAI @@ -44,3 +47,11 @@ def test_openai_stop_error() -> None: llm = OpenAI(stop="3", temperature=0) with pytest.raises(ValueError): llm("write an ordered list of five items", stop=["\n"]) + + +def test_saving_loading_llm(tmp_path: Path) -> None: + """Test saving/loading an OpenAPI LLM.""" + llm = OpenAI(max_tokens=10) + llm.save(file_path=tmp_path / "openai.yaml") + loaded_llm = load_llm(tmp_path / "openai.yaml") + assert loaded_llm == llm diff --git a/tests/integration_tests/llms/utils.py b/tests/integration_tests/llms/utils.py new file mode 100644 index 000000000..6047e979a --- /dev/null +++ b/tests/integration_tests/llms/utils.py @@ -0,0 +1,16 @@ +"""Utils for LLM Tests.""" + +from langchain.llms.base import LLM + + +def assert_llm_equality(llm: LLM, loaded_llm: LLM) -> None: + """Assert LLM Equality for tests.""" + # Check that they are the same type. + assert type(llm) == type(loaded_llm) + # Client field can be session based, so hash is different despite + # all other values being the same, so just assess all other fields + for field in llm.__fields__.keys(): + if field != "client": + val = getattr(llm, field) + new_val = getattr(loaded_llm, field) + assert new_val == val diff --git a/tests/unit_tests/agents/test_agent.py b/tests/unit_tests/agents/test_agent.py index 85853e70b..d407f4b46 100644 --- a/tests/unit_tests/agents/test_agent.py +++ b/tests/unit_tests/agents/test_agent.py @@ -2,17 +2,17 @@ from typing import Any, List, Mapping, Optional +from pydantic import BaseModel + from langchain.agents import Tool, initialize_agent from langchain.llms.base import LLM -class FakeListLLM(LLM): +class FakeListLLM(LLM, BaseModel): """Fake LLM for testing that outputs elements of a list.""" - def __init__(self, responses: List[str]): - """Initialize with list of responses.""" - self.responses = responses - self.i = -1 + responses: List[str] + i: int = -1 def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Increment counter, and then return response in that index.""" @@ -25,6 +25,11 @@ class FakeListLLM(LLM): def _identifying_params(self) -> Mapping[str, Any]: return {} + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "fake_list" + def test_agent_bad_action() -> None: """Test react chain when bad action given.""" @@ -33,7 +38,7 @@ def test_agent_bad_action() -> None: f"I'm turning evil\nAction: {bad_action_name}\nAction Input: misalignment", "Oh well\nAction: Final Answer\nAction Input: curses foiled again", ] - fake_llm = FakeListLLM(responses) + fake_llm = FakeListLLM(responses=responses) tools = [ Tool("Search", lambda x: x, "Useful for searching"), Tool("Lookup", lambda x: x, "Useful for looking up things in a table"), diff --git a/tests/unit_tests/agents/test_react.py b/tests/unit_tests/agents/test_react.py index 2d6c4cca4..917fa7459 100644 --- a/tests/unit_tests/agents/test_react.py +++ b/tests/unit_tests/agents/test_react.py @@ -2,6 +2,8 @@ from typing import Any, List, Mapping, Optional, Union +from pydantic import BaseModel + from langchain.agents.react.base import ReActChain, ReActDocstoreAgent from langchain.agents.tools import Tool from langchain.docstore.base import Docstore @@ -20,13 +22,16 @@ Made in 2022.""" _FAKE_PROMPT = PromptTemplate(input_variables=["input"], template="{input}") -class FakeListLLM(LLM): +class FakeListLLM(LLM, BaseModel): """Fake LLM for testing that outputs elements of a list.""" - def __init__(self, responses: List[str]): - """Initialize with list of responses.""" - self.responses = responses - self.i = -1 + responses: List[str] + i: int = -1 + + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "fake_list" def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """Increment counter, and then return response in that index.""" @@ -50,7 +55,7 @@ class FakeDocstore(Docstore): def test_predict_until_observation_normal() -> None: """Test predict_until_observation when observation is made normally.""" outputs = ["foo\nAction 1: Search[foo]"] - fake_llm = FakeListLLM(outputs) + fake_llm = FakeListLLM(responses=outputs) tools = [ Tool("Search", lambda x: x), Tool("Lookup", lambda x: x), @@ -65,7 +70,7 @@ def test_predict_until_observation_normal() -> None: def test_predict_until_observation_repeat() -> None: """Test when no action is generated initially.""" outputs = ["foo", " Search[foo]"] - fake_llm = FakeListLLM(outputs) + fake_llm = FakeListLLM(responses=outputs) tools = [ Tool("Search", lambda x: x), Tool("Lookup", lambda x: x), @@ -84,7 +89,7 @@ def test_react_chain() -> None: "I should probably lookup\nAction 2: Lookup[made]", "Ah okay now I know the answer\nAction 3: Finish[2022]", ] - fake_llm = FakeListLLM(responses) + fake_llm = FakeListLLM(responses=responses) react_chain = ReActChain(llm=fake_llm, docstore=FakeDocstore()) output = react_chain.run("when was langchain made") assert output == "2022" @@ -97,7 +102,7 @@ def test_react_chain_bad_action() -> None: f"I'm turning evil\nAction 1: {bad_action_name}[langchain]", "Oh well\nAction 2: Finish[curses foiled again]", ] - fake_llm = FakeListLLM(responses) + fake_llm = FakeListLLM(responses=responses) react_chain = ReActChain(llm=fake_llm, docstore=FakeDocstore()) output = react_chain.run("when was langchain made") assert output == "curses foiled again" diff --git a/tests/unit_tests/chains/test_natbot.py b/tests/unit_tests/chains/test_natbot.py index bf4223b29..a9817f712 100644 --- a/tests/unit_tests/chains/test_natbot.py +++ b/tests/unit_tests/chains/test_natbot.py @@ -2,11 +2,13 @@ from typing import Any, List, Mapping, Optional +from pydantic import BaseModel + from langchain.chains.natbot.base import NatBotChain from langchain.llms.base import LLM -class FakeLLM(LLM): +class FakeLLM(LLM, BaseModel): """Fake LLM wrapper for testing purposes.""" def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: @@ -16,6 +18,11 @@ class FakeLLM(LLM): else: return "bar" + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "fake" + @property def _identifying_params(self) -> Mapping[str, Any]: return {} diff --git a/tests/unit_tests/llms/fake_llm.py b/tests/unit_tests/llms/fake_llm.py index 65fd8c793..a3896b6f4 100644 --- a/tests/unit_tests/llms/fake_llm.py +++ b/tests/unit_tests/llms/fake_llm.py @@ -1,20 +1,25 @@ """Fake LLM wrapper for testing purposes.""" from typing import Any, List, Mapping, Optional +from pydantic import BaseModel + from langchain.llms.base import LLM -class FakeLLM(LLM): +class FakeLLM(LLM, BaseModel): """Fake LLM wrapper for testing purposes.""" - def __init__(self, queries: Optional[Mapping] = None): - """Initialize with optional lookup of queries.""" - self._queries = queries + queries: Optional[Mapping] = None + + @property + def _llm_type(self) -> str: + """Return type of llm.""" + return "fake" def __call__(self, prompt: str, stop: Optional[List[str]] = None) -> str: """First try to lookup in queries, else return 'foo' or 'bar'.""" - if self._queries is not None: - return self._queries[prompt] + if self.queries is not None: + return self.queries[prompt] if stop is None: return "foo" else: diff --git a/tests/unit_tests/llms/test_loading.py b/tests/unit_tests/llms/test_loading.py new file mode 100644 index 000000000..e478a0b0b --- /dev/null +++ b/tests/unit_tests/llms/test_loading.py @@ -0,0 +1,15 @@ +"""Test LLM saving and loading functions.""" +from pathlib import Path +from unittest.mock import patch + +from langchain.llms.loading import load_llm +from tests.unit_tests.llms.fake_llm import FakeLLM + + +@patch("langchain.llms.loading.type_to_cls_dict", {"fake": FakeLLM}) +def test_saving_loading_round_trip(tmp_path: Path) -> None: + """Test saving/loading a Fake LLM.""" + fake_llm = FakeLLM() + fake_llm.save(file_path=tmp_path / "fake_llm.yaml") + loaded_llm = load_llm(tmp_path / "fake_llm.yaml") + assert loaded_llm == fake_llm