merge

add embeddings using instructor large endpoint

docs/extras/modules/model_io/models/llms/integrations/octoai.ipynb

@@ -0,0 +1,126 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## OctoAI Compute Service\n",
+    "This example goes over how to use LangChain to interact with `OctoAI` [LLM endpoints](https://octoai.cloud/templates)\n",
+    "## Environment setup\n",
+    "\n",
+    "To run our example app, there are four simple steps to take:\n",
+    "\n",
+    "1. Clone the MPT-7B demo template to your OctoAI account by visiting <https://octoai.cloud/templates/mpt-7b-demo> then clicking \"Clone Template.\" \n",
+    "   1. If you want to use a different LLM model, you can also containerize the model and make a custom OctoAI endpoint yourself, by following [Build a Container from Python](doc:create-custom-endpoints-from-python-code) and [Create a Custom Endpoint from a Container](doc:create-custom-endpoints-from-a-container)\n",
+    "   \n",
+    "2. Paste your Endpoint URL in the code cell below\n",
+    "\n",
+    "3. Get an API Token from [your OctoAI account page](https://octoai.cloud/settings).\n",
+    "   \n",
+    "4. Paste your API key in in the code cell below"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ[\"OCTOAI_API_TOKEN\"] = \"OCTOAI_API_TOKEN\"\n",
+    "os.environ[\"ENDPOINT_URL\"] = \"https://mpt-7b-demo-kk0powt97tmb.octoai.cloud/generate\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from langchain.llms.octoai_endpoint import OctoAIEndpoint\n",
+    "from langchain import PromptTemplate, LLMChain"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "template = \"\"\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n Instruction:\\n{question}\\n Response: \"\"\"\n",
+    "prompt = PromptTemplate(template=template, input_variables=[\"question\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = OctoAIEndpoint(\n",
+    "        model_kwargs={\n",
+    "            \"max_new_tokens\": 200,\n",
+    "            \"temperature\": 0.75,\n",
+    "            \"top_p\": 0.95,\n",
+    "            \"repetition_penalty\": 1,\n",
+    "            \"seed\": None,\n",
+    "            \"stop\": [],\n",
+    "        },\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'\\nLeonardo da Vinci was an Italian polymath and painter regarded by many as one of the greatest painters of all time. He is best known for his masterpieces including Mona Lisa, The Last Supper, and The Virgin of the Rocks. He was a draftsman, sculptor, architect, and one of the most important figures in the history of science. Da Vinci flew gliders, experimented with water turbines and windmills, and invented the catapult and a joystick-type human-powered aircraft control. He may have pioneered helicopters. As a scholar, he was interested in anatomy, geology, botany, engineering, mathematics, and astronomy.\\nOther painters and patrons claimed to be more talented, but Leonardo da Vinci was an incredibly productive artist, sculptor, engineer, anatomist, and scientist.'"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "question = \"Who was leonardo davinci?\"\n",
+    "\n",
+    "llm_chain = LLMChain(prompt=prompt, llm=llm)\n",
+    "\n",
+    "llm_chain.run(question)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "langchain",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.16"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "97697b63fdcee0a640856f91cb41326ad601964008c341809e43189d1cab1047"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

langchain/__init__.py

@@ -117,4 +117,5 @@ __all__ = [
     "PALChain",
     "LlamaCpp",
     "HuggingFaceTextGenInference",
+    "OctoAIEndpoint",
 ]

langchain/embeddings/octoai_embeddings.py

@@ -0,0 +1,92 @@
+"""Module providing a wrapper around OctoAI Compute Service embedding models."""
+
+from typing import Any, Dict, List, Mapping, Optional
+
+from octoai import client
+from pydantic import BaseModel, Extra, Field, root_validator
+
+from langchain.embeddings.base import Embeddings
+from langchain.utils import get_from_dict_or_env
+
+DEFAULT_EMBED_INSTRUCTION = "Represent this input: "
+DEFAULT_QUERY_INSTRUCTION = "Represent the question for retrieving similar documents: "
+
+
+class OctoAIEmbeddings(BaseModel, Embeddings):
+    """
+    Wrapper around OctoAI Compute Service embedding models.
+
+    The environment variable ``OCTOAI_API_TOKEN`` should be set with your API token, or it can be passed
+    as a named parameter to the constructor.
+    """
+
+    endpoint_url: Optional[str] = Field(None, description="Endpoint URL to use.")
+    model_kwargs: Optional[dict] = Field(
+        None, description="Keyword arguments to pass to the model."
+    )
+    octoai_api_token: Optional[str] = Field(None, description="OCTOAI API Token")
+    embed_instruction: str = Field(
+        DEFAULT_EMBED_INSTRUCTION,
+        description="Instruction to use for embedding documents.",
+    )
+    query_instruction: str = Field(
+        DEFAULT_QUERY_INSTRUCTION, description="Instruction to use for embedding query."
+    )
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    @root_validator(allow_reuse=True)
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Ensure that the API key and python package exist in environment."""
+        values["octoai_api_token"] = get_from_dict_or_env(
+            values, "octoai_api_token", "OCTOAI_API_TOKEN"
+        )
+        values["endpoint_url"] = get_from_dict_or_env(
+            values, "endpoint_url", "ENDPOINT_URL"
+        )
+        return values
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Return the identifying parameters."""
+        return {
+            "endpoint_url": self.endpoint_url,
+            "model_kwargs": self.model_kwargs or {},
+        }
+
+    def _compute_embeddings(
+        self, texts: List[str], instruction: str
+    ) -> List[List[float]]:
+        """Common functionality for compute embeddings using a OctoAI instruct model."""
+        embeddings = []
+        octoai_client = client.Client(token=self.octoai_api_token)
+
+        for text in texts:
+            parameter_payload = {
+                "sentence": str([text]),  # for item in text]),
+                "instruction": str([instruction]),  # for item in text]),
+                "parameters": self.model_kwargs or {},
+            }
+
+            try:
+                resp_json = octoai_client.infer(self.endpoint_url, parameter_payload)
+                embedding = resp_json["embeddings"]
+            except Exception as e:
+                raise ValueError(f"Error raised by the inference endpoint: {e}") from e
+
+            embeddings.append(embedding)
+
+        return embeddings
+
+    def embed_documents(self, texts: List[str]) -> List[float]:
+        """Compute document embeddings using an OctoAI instruct model."""
+        texts = list(map(lambda x: x.replace("\n", " "), texts))
+        return self._compute_embeddings(texts, self.embed_instruction)
+
+    def embed_query(self, text: str) -> List[float]:
+        """Compute query embedding using an OctoAI instruct model."""
+        text = text.replace("\n", " ")
+        return self._compute_embeddings([text], self.embed_instruction)

langchain/llms/__init__.py

@@ -31,6 +31,7 @@ from langchain.llms.manifest import ManifestWrapper
 from langchain.llms.modal import Modal
 from langchain.llms.mosaicml import MosaicML
 from langchain.llms.nlpcloud import NLPCloud
+from langchain.llms.octoai_endpoint import OctoAIEndpoint
 from langchain.llms.openai import AzureOpenAI, OpenAI, OpenAIChat
 from langchain.llms.openlm import OpenLM
 from langchain.llms.petals import Petals
@@ -79,6 +80,7 @@ __all__ = [
     "Modal",
     "MosaicML",
     "NLPCloud",
+    "OctoAIEndpoint",
     "OpenAI",
     "OpenAIChat",
     "OpenLM",
@@ -127,6 +129,7 @@ type_to_cls_dict: Dict[str, Type[BaseLLM]] = {
     "modal": Modal,
     "mosaic": MosaicML,
     "nlpcloud": NLPCloud,
+    "octoai_endpoint": OctoAIEndpoint,
     "openai": OpenAI,
     "openlm": OpenLM,
     "petals": Petals,

langchain/llms/octoai_endpoint.py

@@ -0,0 +1,119 @@
+"""Wrapper around OctoAI APIs."""
+from typing import Any, Dict, List, Mapping, Optional
+
+from octoai import client
+from pydantic import Extra, root_validator
+
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+from langchain.llms.base import LLM
+from langchain.llms.utils import enforce_stop_tokens
+from langchain.utils import get_from_dict_or_env
+
+
+class OctoAIEndpoint(LLM):
+    """Wrapper around OctoAI Inference Endpoints.
+    OctoAIEndpoint is a class to interact with OctoAI Compute Service large language model endpoints.
+
+    To use, you should have the ``octoai`` python package installed, and the
+    environment variable ``OCTOAI_API_TOKEN`` set with your API token, or pass
+    it as a named parameter to the constructor.
+
+    Example:
+        .. code-block:: python
+
+            from langchain.llms.octoai_endpoint  import OctoAIEndpoint
+            OctoAIEndpoint(
+                octoai_api_token="octoai-api-key",
+                endpoint_url="https://mpt-7b-demo-kk0powt97tmb.octoai.cloud/generate",
+                model_kwargs={
+                    "max_new_tokens": 200,
+                    "temperature": 0.75,
+                    "top_p": 0.95,
+                    "repetition_penalty": 1,
+                    "seed": None,
+                    "stop": [],
+                },
+            )
+
+    """
+
+    endpoint_url: Optional[str] = None
+    """Endpoint URL to use."""
+
+    model_kwargs: Optional[dict] = None
+    """Key word arguments to pass to the model."""
+
+    octoai_api_token: Optional[str] = None
+    """OCTOAI API Token"""
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    @root_validator(allow_reuse=True)
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        octoai_api_token = get_from_dict_or_env(
+            values, "octoai_api_token", "OCTOAI_API_TOKEN"
+        )
+        values["endpoint_url"] = get_from_dict_or_env(
+            values, "endpoint_url", "ENDPOINT_URL"
+        )
+
+        values["octoai_api_token"] = octoai_api_token
+        return values
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        _model_kwargs = self.model_kwargs or {}
+        return {
+            **{"endpoint_url": self.endpoint_url},
+            **{"model_kwargs": _model_kwargs},
+        }
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "octoai_endpoint"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call out to OctoAI's inference endpoint.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+
+        """
+        _model_kwargs = self.model_kwargs or {}
+
+        # Prepare the payload JSON
+        parameter_payload = {"inputs": prompt, "parameters": _model_kwargs}
+
+        try:
+            # Initialize the OctoAI client
+            octoai_client = client.Client(token=self.octoai_api_token)
+
+            # Send the request using the OctoAI client
+            resp_json = octoai_client.infer(self.endpoint_url, parameter_payload)
+            text = resp_json["generated_text"]
+
+        except Exception as e:
+            # Handle any errors raised by the inference endpoint
+            raise ValueError(f"Error raised by the inference endpoint: {e}") from e
+
+        if stop is not None:
+            # Apply stop tokens when making calls to OctoAI
+            text = enforce_stop_tokens(text, stop)
+
+        return text

tests/integration_tests/embeddings/test_octoai_embeddings.py

@@ -0,0 +1,23 @@
+"""Test octoai embeddings."""
+
+from langchain.embeddings.octoai_embeddings import (
+    OctoAIEmbeddings,
+)
+
+
+def test_octoai_embedding_documents() -> None:
+    """Test octoai embeddings."""
+    documents = ["foo bar"]
+    embedding = OctoAIEmbeddings()
+    output = embedding.embed_documents(documents)
+    assert len(output) == 1
+    assert len(output[0]) == 768
+
+
+def test_octoai_embedding_query() -> None:
+    """Test octoai embeddings."""
+    document = "foo bar"
+    embedding = OctoAIEmbeddings()
+    output = embedding.embed_query(document)
+    assert len(output) == 1
+    assert len(output[0]) == 768

tests/integration_tests/llms/test_octoai_endpoint.py

@@ -0,0 +1,57 @@
+"""Test OctoAI API wrapper."""
+
+import unittest
+from pathlib import Path
+
+import pytest
+
+from langchain.llms.loading import load_llm
+from langchain.llms.octoai_endpoint import OctoAIEndpoint
+from tests.integration_tests.llms.utils import assert_llm_equality
+
+
+def test_octoai_endpoint_text_generation() -> None:
+    """Test valid call to OctoAI text generation model."""
+    llm = OctoAIEndpoint(
+        endpoint_url="https://mpt-7b-demo-kk0powt97tmb.octoai.cloud/generate",
+        model_kwargs={
+            "max_new_tokens": 200,
+            "temperature": 0.75,
+            "top_p": 0.95,
+            "repetition_penalty": 1,
+            "seed": None,
+            "stop": [],
+        },
+    )
+
+    output = llm("Which state is Los Angeles in?")
+    print(output)
+    assert isinstance(output, str)
+
+
+def test_octoai_endpoint_call_error() -> None:
+    """Test valid call to OctoAI that errors."""
+    llm = OctoAIEndpoint(
+        endpoint_url="https://mpt-7b-demo-kk0powt97tmb.octoai.cloud/generate",
+        model_kwargs={"max_new_tokens": -1},
+    )
+    with pytest.raises(ValueError):
+        llm("Which state is Los Angeles in?")
+
+
+def test_saving_loading_endpoint_llm(tmp_path: Path) -> None:
+    """Test saving/loading an OctoAIHub LLM."""
+    llm = OctoAIEndpoint(
+        endpoint_url="https://mpt-7b-demo-kk0powt97tmb.octoai.cloud/generate",
+        model_kwargs={
+            "max_new_tokens": 200,
+            "temperature": 0.75,
+            "top_p": 0.95,
+            "repetition_penalty": 1,
+            "seed": None,
+            "stop": [],
+        },
+    )
+    llm.save(file_path=tmp_path / "octoai.yaml")
+    loaded_llm = load_llm(tmp_path / "octoai.yaml")
+    assert_llm_equality(llm, loaded_llm)