[ { "path": ".gitignore", "content": "env\n.venv\n.ruff_cache\ndist\n*.egg-info\n**/__pycache__" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2023 Matt Rickard\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE." }, { "path": "README.md", "content": "# OpenLM\n\nDrop-in OpenAI-compatible library that can call LLMs from other providers (e.g., HuggingFace, Cohere, and more). \n\n```diff\n1c1\n< import openai\n---\n> import openlm as openai\n\ncompletion = openai.Completion.create(\n model=[\"bloom-560m\", \"cohere.ai/command\"], \n prompt=[\"Hello world!\", \"A second prompt!\"]\n)\nprint(completion)\n```\n### Features\n* Takes in the same parameters as OpenAI's Completion API and returns a similarly structured response. \n* Call models from HuggingFace's inference endpoint API, Cohere.ai, OpenAI, or your custom implementation. \n* Complete multiple prompts on multiple models in the same request. \n* Very small footprint: OpenLM calls the inference APIs directly rather than using multiple SDKs.\n\n\n### Installation\n```bash\npip install openlm\n```\n\n### Examples\n\n- [Import as OpenAI](examples/as_openai.py)\n- [Set up API keys via environment variables or pass a dict](examples/api_keys.py)\n- [Add a custom model or provider](examples/custom_provider.py)\n- [Complete multiple prompts on multiple models](examples/multiplex.py)\n\nOpenLM currently supports the Completion endpoint, but over time will support more standardized endpoints that make sense. \n\n### [Example with Response](examples/multiplex.py)\n\n```python\nimport sys\nfrom pathlib import Path\n\nsys.path.append(str(Path(__file__).resolve().parent.parent))\n\nimport openlm \nimport json\n\ncompletion = openlm.Completion.create(\n model=[\"ada\", \"huggingface.co/gpt2\", \"cohere.ai/command\"],\n prompt=[\"The quick brown fox\", \"Who jumped over the lazy dog?\"],\n max_tokens=15\n)\nprint(json.dumps(completion, indent=4))\n```\n\n```json\n{\n \"id\": \"504cc502-dc27-43e7-bcc3-b62e178c247e\",\n \"object\": \"text_completion\",\n \"created\": 1683583267,\n \"choices\": [\n {\n \"id\": \"c0487ba2-935d-4dec-b191-f7eff962f117\",\n \"model_idx\": 0,\n \"model_name\": \"openai.com/ada\",\n \"index\": 0,\n \"created\": 1683583233,\n \"text\": \" jumps into the much bigger brown bush.\\\" \\\"Alright, people like you can\",\n \"usage\": {\n \"prompt_tokens\": 4,\n \"completion_tokens\": 15,\n \"total_tokens\": 19\n },\n \"extra\": {\n \"id\": \"cmpl-7E3CCSpJHXfx5yB0TaJU9ON7rNYPT\"\n }\n },\n {\n \"id\": \"bab92d11-5ba6-4da2-acca-1f3398a78c3e\",\n \"model_idx\": 0,\n \"model_name\": \"openai.com/ada\",\n \"index\": 1,\n \"created\": 1683583233,\n \"text\": \"\\n\\nIt turns out that saying one's name \\\"Joe\\\" is the\",\n \"usage\": {\n \"prompt_tokens\": 7,\n \"completion_tokens\": 15,\n \"total_tokens\": 22\n },\n \"extra\": {\n \"id\": \"cmpl-7E3CDBbqFy92I2ZbSGoDT5ickAiPD\"\n }\n },\n {\n \"id\": \"be870636-9d9e-4f74-b8bd-d04766072a7b\",\n \"model_idx\": 1,\n \"model_name\": \"huggingface.co/gpt2\",\n \"index\": 0,\n \"created\": 1683583234,\n \"text\": \"The quick brown foxes, and the short, snuggly fox-scented, soft foxes we have in our household\\u2026 all come in two distinct flavours: yellow and orange; and red and white. This mixture is often confused with\"\n },\n {\n \"id\": \"c1abf535-54a9-4b72-8681-d3b4a601da88\",\n \"model_idx\": 1,\n \"model_name\": \"huggingface.co/gpt2\",\n \"index\": 1,\n \"created\": 1683583266,\n \"text\": \"Who jumped over the lazy dog? He probably got it, but there's only so much you do when you lose one.\\n\\nBut I will say for a moment that there's no way this guy might have picked a fight with Donald Trump.\"\n },\n {\n \"id\": \"08e8c351-236a-4497-98f3-488cdc0b6b6a\",\n \"model_idx\": 2,\n \"model_name\": \"cohere.ai/command\",\n \"index\": 0,\n \"created\": 1683583267,\n \"text\": \"\\njumps over the lazy dog.\",\n \"extra\": {\n \"request_id\": \"0bbb28c0-eb3d-4614-b4d9-1eca88c361ca\",\n \"generation_id\": \"5288dd6f-3ecf-475b-b909-0b226be6a193\"\n }\n },\n {\n \"id\": \"49ce51e6-9a18-4093-957f-54a1557c8829\",\n \"model_idx\": 2,\n \"model_name\": \"cohere.ai/command\",\n \"index\": 1,\n \"created\": 1683583267,\n \"text\": \"\\nThe quick brown fox.\",\n \"extra\": {\n \"request_id\": \"ab5d5e03-22a1-42cd-85b2-9b9704c79304\",\n \"generation_id\": \"60493966-abf6-483c-9c47-2ea5c5eeb855\"\n }\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 11,\n \"completion_tokens\": 30,\n \"total_tokens\": 41\n }\n}\n```\n\n### Other Languages\n[r2d4/llm.ts](https://github.com/r2d4/llm.ts) is a TypeScript library that has a similar API that sits on top of multiple language models.\n\n### Roadmap\n- [ ] Streaming API\n- [ ] Embeddings API\n\n### Contributing\nContributions are welcome! Please open an issue or submit a PR.\n\n### License\n[MIT](LICENSE)\n\n" }, { "path": "examples/api_keys.py", "content": "import sys\nfrom pathlib import Path\n\nsys.path.append(str(Path(__file__).resolve().parent.parent))\n\nimport openlm\nimport json\n\ncompletion = openlm.Completion.create(\n model=[\"ada\", \"distilgpt2\", \"huggingface.co/\"], \n prompt=\"Hello world\",\n api_keys={\n 'huggingface.co': 'YOUR_API_KEY', # or os.environ[\"HF_API_TOKEN\"]\n 'cohere.ai': 'YOUR_API_KEY', # or os.environ[\"COHERE_API_KEY\"]\n 'openai.com': 'YOUR_API_KEY' # or os.environ[\"OPENAI_API_KEY\"]\n },\n)\n\nprint(json.dumps(completion, indent=4))\n" }, { "path": "examples/as_openai.py", "content": "import sys\nfrom pathlib import Path\n\nsys.path.append(str(Path(__file__).resolve().parent.parent))\n\nimport openlm as openai\nimport json\n\ncompletion = openai.Completion.create(\n model=\"ada\", \n prompt=\"Hello world\"\n)\nprint(json.dumps(completion, indent=4))\n\n'''\n{\n \"id\": \"3890d5c3-e6c4-4222-b77d-40a65f1d032b\",\n \"object\": \"text_completion\",\n \"created\": 1683583320,\n \"choices\": [\n {\n \"id\": \"660d576d-8c04-420f-b410-146729c8fc8a\",\n \"model_idx\": 0,\n \"model_name\": \"openai.com/ada\",\n \"index\": 0,\n \"created\": 1683583320,\n \"text\": \". Details?), website modding, something that makes no sense, and just design\",\n \"usage\": {\n \"prompt_tokens\": 2,\n \"completion_tokens\": 16,\n \"total_tokens\": 18\n },\n \"extra\": {\n \"id\": \"cmpl-7E3DcN1o6Axv1JIXxHNt11Q0wA1Is\"\n }\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 2,\n \"completion_tokens\": 16,\n \"total_tokens\": 18\n }\n}\n'''" }, { "path": "examples/custom_provider.py", "content": "import sys\nfrom pathlib import Path\n\nsys.path.append(str(Path(__file__).resolve().parent.parent))\n\nimport openlm\nimport json\nfrom typing import Any, Dict, List, Optional, Union\n\nclass CustomModel(openlm.BaseModel):\n def create_completion(self, model: Union[str, List[str]], prompt: Union[str, List[str]],\n suffix: Optional[str] = None,\n max_tokens: Optional[int] = None,\n temperature: Optional[float] = None,\n top_p: Optional[float] = None,\n n: Optional[int] = None,\n stream: Optional[bool] = None,\n logprobs: Optional[int] = None,\n echo: Optional[bool] = None,\n stop: Optional[Union[str, List[str]]] = None,\n presence_penalty: Optional[float] = None,\n frequency_penalty: Optional[float] = None,\n best_of: Optional[int] = None,\n logit_bias: Optional[Dict[str, float]] = None,\n user: Optional[str] = None) -> Dict[str, Any]:\n # completions should return a dictionary with the following keys:\n return {\n # Required keys:\n \"text\": \"Hello world!\"\n\n ## Optional keys:\n # ,'extra': {\n # 'key': 'value'\n # },\n # 'usage': {\n # 'prompt_tokens': 0,\n # 'completion_tokens': 0,\n # 'total_tokens': 0,\n # }\n }\n \n def list_models(self) -> Dict[str, Any]:\n # list of model names that can be used with this provider\n return [\"your_model_name\"]\n \n def namespace(self) -> str:\n # A namespace prevents name collisions between models from different providers.\n # You will be able to reference your model both as:\n # your_namespace/your_model_name or your_model_name\n return \"your_namespace\"\n \nopenlm.Completion.register(CustomModel())\n\n# Now you can use your custom model in the same way as the built-in models:\ncompletion = openlm.Completion.create(\n model=\"your_model_name\",\n prompt=\"Hello world\"\n)\n\nprint(json.dumps(completion, indent=4))\n\n'''\n{\n \"id\": \"12bf5515-e2cc-463d-b120-c21c911364f9\",\n \"object\": \"text_completion\",\n \"created\": 1683583298,\n \"choices\": [\n {\n \"id\": \"2dde9e4e-17c3-4d92-be6f-285fb9a96935\",\n \"model_idx\": 0,\n \"model_name\": \"your_namespace/your_model_name\",\n \"index\": 0,\n \"created\": 1683583298,\n \"text\": \"Hello world!\"\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 0,\n \"completion_tokens\": 0,\n \"total_tokens\": 0\n }\n}\n'''" }, { "path": "examples/multiplex.py", "content": "import sys\nfrom pathlib import Path\n\nsys.path.append(str(Path(__file__).resolve().parent.parent))\n\nimport openlm \nimport json\n\ncompletion = openlm.Completion.create(\n model=[\"ada\", \"huggingface.co/gpt2\", \"cohere.ai/command\"],\n prompt=[\"The quick brown fox\", \"Who jumped over the lazy dog?\"],\n max_tokens=15\n)\nprint(json.dumps(completion, indent=4))\n\n'''\n{\n \"id\": \"504cc502-dc27-43e7-bcc3-b62e178c247e\",\n \"object\": \"text_completion\",\n \"created\": 1683583267,\n \"choices\": [\n {\n \"id\": \"c0487ba2-935d-4dec-b191-f7eff962f117\",\n \"model_idx\": 0,\n \"model_name\": \"openai.com/ada\",\n \"index\": 0,\n \"created\": 1683583233,\n \"text\": \" jumps into the much bigger brown bush.\\\" \\\"Alright, people like you can\",\n \"usage\": {\n \"prompt_tokens\": 4,\n \"completion_tokens\": 15,\n \"total_tokens\": 19\n },\n \"extra\": {\n \"id\": \"cmpl-7E3CCSpJHXfx5yB0TaJU9ON7rNYPT\"\n }\n },\n {\n \"id\": \"bab92d11-5ba6-4da2-acca-1f3398a78c3e\",\n \"model_idx\": 0,\n \"model_name\": \"openai.com/ada\",\n \"index\": 1,\n \"created\": 1683583233,\n \"text\": \"\\n\\nIt turns out that saying one's name \\\"Joe\\\" is the\",\n \"usage\": {\n \"prompt_tokens\": 7,\n \"completion_tokens\": 15,\n \"total_tokens\": 22\n },\n \"extra\": {\n \"id\": \"cmpl-7E3CDBbqFy92I2ZbSGoDT5ickAiPD\"\n }\n },\n {\n \"id\": \"be870636-9d9e-4f74-b8bd-d04766072a7b\",\n \"model_idx\": 1,\n \"model_name\": \"huggingface.co/gpt2\",\n \"index\": 0,\n \"created\": 1683583234,\n \"text\": \"The quick brown foxes, and the short, snuggly fox-scented, soft foxes we have in our household\\u2026 all come in two distinct flavours: yellow and orange; and red and white. This mixture is often confused with\"\n },\n {\n \"id\": \"c1abf535-54a9-4b72-8681-d3b4a601da88\",\n \"model_idx\": 1,\n \"model_name\": \"huggingface.co/gpt2\",\n \"index\": 1,\n \"created\": 1683583266,\n \"text\": \"Who jumped over the lazy dog? He probably got it, but there's only so much you do when you lose one.\\n\\nBut I will say for a moment that there's no way this guy might have picked a fight with Donald Trump.\"\n },\n {\n \"id\": \"08e8c351-236a-4497-98f3-488cdc0b6b6a\",\n \"model_idx\": 2,\n \"model_name\": \"cohere.ai/command\",\n \"index\": 0,\n \"created\": 1683583267,\n \"text\": \"\\njumps over the lazy dog.\",\n \"extra\": {\n \"request_id\": \"0bbb28c0-eb3d-4614-b4d9-1eca88c361ca\",\n \"generation_id\": \"5288dd6f-3ecf-475b-b909-0b226be6a193\"\n }\n },\n {\n \"id\": \"49ce51e6-9a18-4093-957f-54a1557c8829\",\n \"model_idx\": 2,\n \"model_name\": \"cohere.ai/command\",\n \"index\": 1,\n \"created\": 1683583267,\n \"text\": \"\\nThe quick brown fox.\",\n \"extra\": {\n \"request_id\": \"ab5d5e03-22a1-42cd-85b2-9b9704c79304\",\n \"generation_id\": \"60493966-abf6-483c-9c47-2ea5c5eeb855\"\n }\n }\n ],\n \"usage\": {\n \"prompt_tokens\": 11,\n \"completion_tokens\": 30,\n \"total_tokens\": 41\n }\n}\n'''" }, { "path": "openlm/__init__.py", "content": "from openlm.openlm import Completion\nfrom openlm.llm.base import BaseModel\n\n# For backwards compatibility with OpenAI\napi_key = None" }, { "path": "openlm/llm/__init__.py", "content": "from .base import BaseModel\nfrom .openai import OpenAI\nfrom .huggingface import Huggingface\nfrom .cohere import Cohere\n" }, { "path": "openlm/llm/base.py", "content": "import abc\nfrom typing import Any, Dict, List, Optional, Union\n\n\nclass BaseCompletion(metaclass=abc.ABCMeta):\n @abc.abstractmethod\n def create_completion(self, model: Union[str, List[str]], prompt: Union[str, List[str]],\n suffix: Optional[str] = None,\n max_tokens: Optional[int] = None,\n temperature: Optional[float] = None,\n top_p: Optional[float] = None,\n n: Optional[int] = None,\n stream: Optional[bool] = None,\n logprobs: Optional[int] = None,\n echo: Optional[bool] = None,\n stop: Optional[Union[str, List[str]]] = None,\n presence_penalty: Optional[float] = None,\n frequency_penalty: Optional[float] = None,\n best_of: Optional[int] = None,\n logit_bias: Optional[Dict[str, float]] = None,\n user: Optional[str] = None) -> Dict[str, Any]:\n raise NotImplementedError\n\n\nclass BaseModel(BaseCompletion, metaclass=abc.ABCMeta):\n @abc.abstractmethod\n def list_models(self) -> Dict[str, Any]:\n raise NotImplementedError\n \n @abc.abstractmethod\n def namespace(self) -> str:\n raise NotImplementedError " }, { "path": "openlm/llm/cohere.py", "content": "\n\nfrom openlm.llm.base import BaseModel\nimport os\nfrom typing import Any, Dict, List, Optional, Union\nimport json\nimport requests\n\ncohere_models = [\n 'command',\n 'command-nightly',\n 'command-light',\n 'command-light-nightly',\n]\n\nclass Cohere(BaseModel):\n def __init__(self,\n api_key = os.environ.get(\"COHERE_API_KEY\"),\n model_list = cohere_models,\n namespace = 'cohere.ai',\n base_url = 'https://api.cohere.ai/v1/generate'):\n self.api_key = api_key\n self.model_list = model_list\n self._namespace = namespace\n self.base_url = base_url\n\n def list_models(self):\n return self.model_list\n \n def namespace(self):\n return self._namespace\n \n def create_completion(self, model: Union[str, List[str]], prompt: Union[str, List[str]],\n suffix: Optional[str] = None,\n max_tokens: Optional[int] = None,\n temperature: Optional[float] = None,\n top_p: Optional[float] = None,\n n: Optional[int] = None,\n stream: Optional[bool] = None,\n logprobs: Optional[int] = None,\n echo: Optional[bool] = None,\n stop: Optional[Union[str, List[str]]] = None,\n presence_penalty: Optional[float] = None,\n frequency_penalty: Optional[float] = None,\n best_of: Optional[int] = None,\n logit_bias: Optional[Dict[str, float]] = None,\n user: Optional[str] = None) -> Dict[str, Any]:\n \n headers = {'Content-Type': 'application/json',\n 'Authorization': f'Bearer {self.api_key}'}\n payload = {\n 'prompt': prompt,\n 'model': model,\n 'max_tokens': max_tokens,\n 'temperature': temperature,\n 'p': top_p,\n 'frequency_penalty': frequency_penalty,\n 'presence_penalty': presence_penalty,\n 'stop_sequences': stop,\n }\n payload_str = json.dumps({k: v for k, v in payload.items() if v is not None})\n resp = requests.post(self.base_url, headers=headers, data=payload_str)\n if resp.status_code != 200:\n raise ValueError(resp.status_code, resp.text)\n return self._convert_response(resp.json())\n\n def _convert_request(req):\n return {\n 'prompt': req.prompt,\n 'top_p': req.top_p,\n 'temperature': req.temperature,\n 'max_new_tokens': req.max_tokens,\n }\n \n def _convert_response(self, resp):\n return {\n 'text': resp['generations'][0]['text'],\n 'extra': {\n 'request_id': resp['id'],\n 'generation_id': resp['generations'][0]['id'],\n }\n }\n" }, { "path": "openlm/llm/huggingface.py", "content": "\n\nfrom openlm.llm.base import BaseModel\nimport os\nfrom typing import Any, Dict, List, Optional, Union\nimport json\nimport requests\n\nhf_models = [\n 'gpt2',\n 'distilgpt2',\n 'gpt2-large',\n 'gpt2-medium',\n 'gpt2-xl',\n\n 'bigscience/bloom-560m',\n 'bigscience/bloom-1b',\n 'bigscience/bloom-3b',\n 'bigscience/bloom-7b1',\n\n 'decapoda-research/llama-7b-hf',\n 'decapoda-research/llama-13b-hf',\n 'decapoda-research/llama-30b-hf',\n 'decapoda-research/llama-65b-hf',\n \n 'EleutherAI/gpt-j-6B',\n 'EleutherAI/gpt-j-2.7B',\n\n 'EleutherAI/gpt-neo-125M',\n 'EleutherAI/gpt-neo-1.3B',\n 'EleutherAI/gpt-neox-20B',\n\n 'EleutherAI/pythia-160m',\n 'EleutherAI/pythia-70m',\n 'EleutherAI/pythia-12b',\n \n 'cerebras/Cerebras-GPT-111M',\n 'cerebras/Cerebras-GPT-1.3B',\n 'cerebras/Cerebras-GPT-2.7B',\n \n 'bigcode/santacoder',\n \n 'Salesforce/codegen-350M-multi',\n 'Salesforce/codegen-2b-multi',\n \n 'stabilityai/stablelm-tuned-alpha-3b',\n 'stabilityai/stablelm-tuned-alpha-7b',\n\n 'facebook/opt-125m',\n 'facebook/opt-350m',\n 'facebook/opt-1.3b',\n 'facebook/opt-2.7b',\n 'facebook/opt-6.7b',\n 'facebook/opt-13b',\n 'facebook/opt-30b',\n\n 'mosaicml/mpt-7b',\n 'mosaicml/mpt-7b-instruct',\n \n 'databricks/dolly-v2-7b',\n 'databricks/dolly-v2-12b',\n\n\n]\n\nclass Huggingface(BaseModel):\n def __init__(self,\n api_key = os.environ.get(\"HF_API_TOKEN\"),\n model_list = hf_models,\n namespace = 'huggingface.co',\n base_url = 'https://api-inference.huggingface.co/models'):\n self.api_key = api_key\n self.model_list = model_list\n self._namespace = namespace\n self.base_url = base_url\n\n def list_models(self):\n return self.model_list\n \n def namespace(self):\n return self._namespace\n \n def create_completion(self, model: Union[str, List[str]], prompt: Union[str, List[str]],\n suffix: Optional[str] = None,\n max_tokens: Optional[int] = None,\n temperature: Optional[float] = None,\n top_p: Optional[float] = None,\n n: Optional[int] = None,\n stream: Optional[bool] = None,\n logprobs: Optional[int] = None,\n echo: Optional[bool] = None,\n stop: Optional[Union[str, List[str]]] = None,\n presence_penalty: Optional[float] = None,\n frequency_penalty: Optional[float] = None,\n best_of: Optional[int] = None,\n logit_bias: Optional[Dict[str, float]] = None,\n user: Optional[str] = None) -> Dict[str, Any]:\n \n headers = {'Content-Type': 'application/json',\n 'Authorization': f'Bearer {self.api_key}'}\n payload = {\n 'inputs': prompt,\n 'top_p': top_p,\n 'temperature': temperature,\n 'max_new_tokens': max_tokens,\n }\n payload_str = json.dumps({k: v for k, v in payload.items() if v is not None})\n resp = requests.post(self.base_url + '/' + model, headers=headers, data=payload_str)\n if resp.status_code != 200:\n raise ValueError(resp.status_code, resp.text)\n return self._convert_response(resp.json())\n\n def _convert_request(req):\n return {\n 'prompt': req.prompt,\n 'top_p': req.top_p,\n 'temperature': req.temperature,\n 'max_new_tokens': req.max_tokens,\n }\n \n def _convert_response(self, resp):\n return {\n 'text': resp[0]['generated_text'],\n }\n" }, { "path": "openlm/llm/openai.py", "content": "import json\nimport os\nfrom typing import Any, Dict, List, Optional, Union\n\nimport requests\n\nfrom openlm.llm.base import BaseModel\n\nopenai_models = [\n 'text-davinci-003', \n 'text-davinci-002', \n 'text-curie-001', \n 'text-babbage-001', \n 'text-ada-001',\n \n # aliases\n 'ada',\n 'babbage',\n 'curie',\n 'davinci',\n ]\n\nclass OpenAI(BaseModel):\n def __init__(self,\n api_key = os.environ.get(\"OPENAI_API_KEY\"), \n model_list = openai_models, \n namespace = 'openai.com', \n base_url = 'https://api.openai.com/v1/completions'):\n \n if api_key is None:\n raise ValueError(\"OPENAI_API_KEY is not set or passed as an argument\")\n \n self.api_key = api_key\n self.model_list = model_list\n self._namespace = namespace\n self.base_url = base_url\n\n def create_completion(self, model: Union[str, List[str]], prompt: Union[str, List[str]],\n suffix: Optional[str] = None,\n max_tokens: Optional[int] = None,\n temperature: Optional[float] = None,\n top_p: Optional[float] = None,\n n: Optional[int] = None,\n stream: Optional[bool] = None,\n logprobs: Optional[int] = None,\n echo: Optional[bool] = None,\n stop: Optional[Union[str, List[str]]] = None,\n presence_penalty: Optional[float] = None,\n frequency_penalty: Optional[float] = None,\n best_of: Optional[int] = None,\n logit_bias: Optional[Dict[str, float]] = None,\n user: Optional[str] = None) -> Dict[str, Any]:\n headers = {'Content-Type': 'application/json',\n 'Authorization': f'Bearer {self.api_key}'}\n payload = {\n 'model': model,\n 'prompt': prompt,\n 'suffix': suffix,\n 'max_tokens': max_tokens,\n 'temperature': temperature,\n 'top_p': top_p,\n 'n': n,\n 'stream': stream,\n 'logprobs': logprobs,\n 'echo': echo,\n 'stop': stop,\n 'presence_penalty': presence_penalty,\n 'frequency_penalty': frequency_penalty,\n 'best_of': best_of,\n 'logit_bias': logit_bias,\n 'user': user\n }\n\n payload_str = json.dumps({k: v for k, v in payload.items() if v is not None})\n resp = requests.post(self.base_url, headers=headers, data=payload_str).json()\n if 'error' in resp:\n raise ValueError(resp['error'])\n return self._convert_response(resp)\n\n def _convert_response(self, response: Dict[str, Any]) -> Dict[str, Any]:\n return {\n 'text': response['choices'][0]['text'],\n 'extra': {\n 'id': response['id'],\n },\n 'usage': response['usage'],\n }\n\n def list_models(self):\n return self.model_list\n \n def namespace(self):\n return self._namespace" }, { "path": "openlm/openlm.py", "content": "\nfrom typing import Any, Dict, List, Optional, Union\nimport uuid\nfrom openlm.llm import BaseModel, OpenAI, Huggingface, Cohere\nimport time\nimport openlm\nfrom concurrent.futures import ThreadPoolExecutor\n\nclass Completion():\n \"\"\"\n OpenAI-compatible completion API\n \"\"\"\n models = {}\n aliases = {}\n\n @classmethod\n def create(cls, model: Union[str, List[str]], prompt: Union[str, List[str]],\n suffix: Optional[str] = None,\n max_tokens: Optional[int] = None,\n temperature: Optional[float] = None,\n top_p: Optional[float] = None,\n n: Optional[int] = None,\n stream: Optional[bool] = None,\n logprobs: Optional[int] = None,\n echo: Optional[bool] = None,\n stop: Optional[Union[str, List[str]]] = None,\n presence_penalty: Optional[float] = None,\n frequency_penalty: Optional[float] = None,\n best_of: Optional[int] = None,\n logit_bias: Optional[Dict[str, float]] = None,\n user: Optional[str] = None,\n api_keys: Optional[Dict[str, str]] = None,\n request_timeout=0) -> Dict[str, Any]:\n \"\"\"\n Creates a completion request for the OpenAI API.\n\n :param model: The ID(s) of the model to use.\n :param prompt: The prompt(s) to generate completions for.\n :param suffix: A string to append to the completion(s).\n :param max_tokens: The maximum number of tokens to generate in the completion(s).\n :param temperature: The sampling temperature to use.\n :param top_p: The nucleus sampling probability to use.\n :param n: The number of completions to generate.\n :param stream: Whether to stream back partial progress updates.\n :param logprobs: The number of log probabilities to generate per token.\n :param echo: Whether to include the prompt(s) in the completion(s).\n :param stop: The stop sequence(s) to use.\n :param presence_penalty: The presence penalty to use.\n :param frequency_penalty: The frequency penalty to use.\n :param best_of: The number of completions to generate and return the best of.\n :param logit_bias: A dictionary of token IDs and bias values to use.\n :param user: The ID of the user making the request.\n :return: A dictionary containing the completion response.\n \"\"\"\n cls.register_default()\n if isinstance(model, str):\n model = [model]\n\n if isinstance(prompt, str):\n prompt = [prompt]\n\n # Create a list of tuples, each containing all the parameters for a call to _generate_completion\n args = [(m, p, suffix, max_tokens, temperature, top_p, n, stream, logprobs, echo, stop, presence_penalty, frequency_penalty, best_of, logit_bias, user) \n for m in model for p in prompt]\n\n total_usage = {\n 'prompt_tokens': 0,\n 'completion_tokens': 0,\n 'total_tokens': 0\n }\n\n # Use a ThreadPoolExecutor to run _generate_completion in parallel for each set of parameters\n with ThreadPoolExecutor() as executor:\n choices = list(executor.map(lambda params: cls._generate_completion(*params), args))\n\n # Sum up the usage from all choices\n for choice in choices:\n if 'usage' in choice:\n total_usage['prompt_tokens'] += choice['usage']['prompt_tokens']\n total_usage['completion_tokens'] += choice['usage']['completion_tokens']\n total_usage['total_tokens'] += choice['usage']['total_tokens']\n\n return {\n \"id\": str(uuid.uuid4()),\n \"object\": \"text_completion\",\n \"created\": int(time.time()),\n \"choices\": choices,\n \"usage\": total_usage,\n }\n \n @classmethod\n def _generate_completion(cls, model, prompt, suffix, max_tokens, temperature, top_p, n, stream, logprobs, echo, stop, presence_penalty, frequency_penalty, best_of, logit_bias, user):\n \"\"\"\n Function to generate a single completion. This will be used in parallel execution.\n \"\"\"\n if model not in cls.aliases:\n raise ValueError(f\"Model {model} not found. OpenLM currently supports the following models:\\n{cls._pretty_list_models()}\")\n fqn = cls.aliases[model]\n try:\n ret = cls.models[fqn].create_completion(\n model=fqn[len(cls.models[fqn].namespace())+1:],\n prompt=prompt, \n suffix=suffix, \n max_tokens=max_tokens, \n temperature=temperature, \n top_p=top_p, \n n=n, \n stream=stream, \n logprobs=logprobs, \n echo=echo, \n stop=stop, \n presence_penalty=presence_penalty, \n frequency_penalty=frequency_penalty, \n best_of=best_of, \n logit_bias=logit_bias, \n user=user)\n except Exception as e:\n ret = {\n 'error': f\"Error: {e}\"\n }\n choice = {\n \"id\": str(uuid.uuid4()),\n \"model_name\": fqn,\n 'created': int(time.time()),\n }\n if 'error' in ret:\n choice['error'] = ret['error']\n if 'text' in ret:\n choice['text'] = ret['text']\n if 'usage' in ret:\n choice['usage'] = ret['usage']\n if 'extra' in ret:\n choice['extra'] = ret['extra']\n return choice\n \n @classmethod\n def register(cls, providers: BaseModel | List[BaseModel]):\n if not isinstance(providers, list):\n providers = [providers]\n for provider in providers:\n for model in provider.list_models():\n fqn = provider.namespace() + '/' + model\n cls.models[fqn] = provider\n cls.aliases[model] = fqn\n cls.aliases[fqn] = fqn\n if '/' in model:\n cls.aliases[model.split('/')[1]] = fqn\n\n @classmethod\n def register_default(cls, api_keys: Optional[Dict[str, str]] = None):\n if openlm.api_key:\n cls.register(OpenAI(api_key=openlm.api_key))\n else:\n if api_keys and api_keys['openai.com'] is not None:\n cls.register(OpenAI(api_key=api_keys['openai.com']))\n else:\n cls.register(OpenAI())\n if api_keys and api_keys['huggingface.co'] is not None:\n cls.register(Huggingface(api_key=api_keys['huggingface.co']))\n else:\n cls.register(Huggingface())\n if api_keys and api_keys['cohere.ai'] is not None:\n cls.register(Cohere(api_key=api_keys['cohere.ai']))\n else:\n cls.register(Cohere())\n\n @classmethod\n def list_models(cls) -> List[str]:\n reverse_alias = {}\n for key, value in cls.aliases.items():\n # If the value is not in the reverse dictionary, create an empty array for it\n if value not in reverse_alias:\n reverse_alias[value] = []\n # Append the key to the array for the value in the reverse dictionary\n reverse_alias[value].append(key)\n\n return reverse_alias\n \n @classmethod\n def _pretty_list_models(cls):\n ret = \"\"\n for key, value in cls.list_models().items():\n ret += f\"-> {value} \\n\"\n return ret\n\n\n" }, { "path": "pyproject.toml", "content": "[tool.poetry]\nname = \"openlm\"\nversion = \"0.0.4\"\ndescription = \"Drop-in OpenAI-compatible that can call LLMs from other providers\"\nauthors = [\"Matt Rickard \"]\nmaintainers = [\"Matt Rickard \"]\nreadme = \"README.md\"\nlicense = \"MIT\"\nkeywords = [\n \"llm\",\n \"ai\",\n \"prompt\",\n \"large language models\",\n \"gpt-3\",\n \"chatgpt\",\n]\nclassifiers = [\n \"Development Status :: 3 - Alpha\",\n \"Intended Audience :: Developers\",\n \"Intended Audience :: Science/Research\",\n \"License :: OSI Approved :: MIT License\",\n \"Programming Language :: Python :: 3\",\n \"Programming Language :: Python :: 3.7\",\n \"Programming Language :: Python :: 3.8\",\n \"Programming Language :: Python :: 3.9\",\n \"Programming Language :: Python :: 3 :: Only\",\n \"Topic :: Scientific/Engineering :: Artificial Intelligence\",\n \"Topic :: Software Development :: Libraries :: Python Modules\",\n \"Topic :: Text Processing :: Linguistic\",\n]\nurls = { repository = \"https://github.com/r2d4/openlm\" }\n\n\n[tool.poetry.dependencies]\npython = \">=3.8.1,<4.0\"\nrequests = \"^2\"\n\n\n[build-system]\nrequires = [\"poetry-core\"]\nbuild-backend = \"poetry.core.masonry.api\"\n" } ]