From 5ddb2b22af1159c544eab2408c8bb8298d32e5d0 Mon Sep 17 00:00:00 2001 From: Guofeng Yi <1548222878@qq.com> Date: Wed, 26 Jun 2024 00:41:07 +0800 Subject: [PATCH] Integrate Yi model (#14353) * first update * modify * support yi llm * fix url error * rm colab link * pants and lint * pants and lint --------- Co-authored-by: Andrei Fajardo --- docs/docs/examples/llm/yi.ipynb | 232 ++++++++++++++++++ .../llms/llama-index-llms-yi/.gitignore | 153 ++++++++++++ .../llms/llama-index-llms-yi/BUILD | 3 + .../llms/llama-index-llms-yi/Makefile | 17 ++ .../llms/llama-index-llms-yi/README.md | 1 + .../llama_index/llms/yi/BUILD | 1 + .../llama_index/llms/yi/__init__.py | 4 + .../llama_index/llms/yi/base.py | 198 +++++++++++++++ .../llms/llama-index-llms-yi/pyproject.toml | 58 +++++ .../llms/llama-index-llms-yi/tests/BUILD | 1 + .../llama-index-llms-yi/tests/__init__.py | 0 .../llama-index-llms-yi/tests/test_llms_yi.py | 7 + 12 files changed, 675 insertions(+) create mode 100644 docs/docs/examples/llm/yi.ipynb create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/.gitignore create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/BUILD create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/Makefile create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/README.md create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/BUILD create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/__init__.py create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/base.py create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/pyproject.toml create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/tests/BUILD create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/tests/__init__.py create mode 100644 llama-index-integrations/llms/llama-index-llms-yi/tests/test_llms_yi.py diff --git a/docs/docs/examples/llm/yi.ipynb b/docs/docs/examples/llm/yi.ipynb new file mode 100644 index 0000000000000..2099e86146aec --- /dev/null +++ b/docs/docs/examples/llm/yi.ipynb @@ -0,0 +1,232 @@ +{ + "cells": [ + { + "attachments": {}, + "cell_type": "markdown", + "id": "2e33dced-e587-4397-81b3-d6606aa1738a", + "metadata": {}, + "source": [ + "# 01.AI LLM\n", + "\n", + "This notebook shows how to use `Yi` series LLMs.\n", + "\n", + "Visit https://platform.01.ai/ and sign up to get an API key." + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "5863dde9-84a0-4c33-ad52-cc767442f63f", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "markdown", + "id": "833bdb2b", + "metadata": {}, + "source": [ + "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4aff387e", + "metadata": {}, + "outputs": [], + "source": [ + "%pip install llama-index-llms-yi" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9bbbc106", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install llama-index" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ad297f19-998f-4485-aa2f-d67020058b7d", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.llms.yi import Yi" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "152ced37-9a42-47be-9a39-4218521f5e72", + "metadata": {}, + "outputs": [], + "source": [ + "# set api key in env or in llm\n", + "# import os\n", + "# os.environ[\"YI_API_KEY\"] = \"your api key\"\n", + "\n", + "llm = Yi(model=\"yi-large\", api_key=\"your api key\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d61b10bb-e911-47fb-8e84-19828cf224be", + "metadata": {}, + "outputs": [], + "source": [ + "resp = llm.complete(\"Who is Paul Graham?\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3bd14f4e-c245-4384-a471-97e4ddfcb40e", + "metadata": {}, + "outputs": [], + "source": [ + "print(resp)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "3ba9503c-b440-43c6-a50c-676c79993813", + "metadata": {}, + "source": [ + "#### Call `chat` with a list of messages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ee8a4a55-5680-4dc6-a44c-fc8ad7892f80", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.llms import ChatMessage\n", + "\n", + "messages = [\n", + " ChatMessage(\n", + " role=\"system\", content=\"You are a pirate with a colorful personality\"\n", + " ),\n", + " ChatMessage(role=\"user\", content=\"What is your name\"),\n", + "]\n", + "resp = llm.chat(messages)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2a9bfe53-d15b-4e75-9d91-8c5d024f4eda", + "metadata": {}, + "outputs": [], + "source": [ + "print(resp)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "25ad1b00-28fc-4bcd-96c4-d5b35605721a", + "metadata": {}, + "source": [ + "### Streaming" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "13c641fa-345a-4dce-87c5-ab1f6dcf4757", + "metadata": {}, + "source": [ + "Using `stream_complete` endpoint " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "06da1ef1-2f6b-497c-847b-62dd2df11491", + "metadata": {}, + "outputs": [], + "source": [ + "response = llm.stream_complete(\"Who is Paul Graham?\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1b851def-5160-46e5-a30c-5a3ef2356b79", + "metadata": {}, + "outputs": [], + "source": [ + "for r in response:\n", + " print(r.delta, end=\"\")" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "id": "ca52051d-6b28-49d7-98f5-82e266a1c7a6", + "metadata": {}, + "source": [ + "Using `stream_chat` endpoint" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fe553190-52a9-436d-84ae-4dd99a1808f4", + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core.llms import ChatMessage\n", + "\n", + "messages = [\n", + " ChatMessage(\n", + " role=\"system\", content=\"You are a pirate with a colorful personality\"\n", + " ),\n", + " ChatMessage(role=\"user\", content=\"What is your name\"),\n", + "]\n", + "resp = llm.stream_chat(messages)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "154c503c-f893-4b6b-8a65-a9a27b636046", + "metadata": {}, + "outputs": [], + "source": [ + "for r in resp:\n", + " print(r.delta, end=\"\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/llama-index-integrations/llms/llama-index-llms-yi/.gitignore b/llama-index-integrations/llms/llama-index-llms-yi/.gitignore new file mode 100644 index 0000000000000..990c18de22908 --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/.gitignore @@ -0,0 +1,153 @@ +llama_index/_static +.DS_Store +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +bin/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +etc/ +include/ +lib/ +lib64/ +parts/ +sdist/ +share/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +.ruff_cache + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints +notebooks/ + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ +pyvenv.cfg + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# Jetbrains +.idea +modules/ +*.swp + +# VsCode +.vscode + +# pipenv +Pipfile +Pipfile.lock + +# pyright +pyrightconfig.json diff --git a/llama-index-integrations/llms/llama-index-llms-yi/BUILD b/llama-index-integrations/llms/llama-index-llms-yi/BUILD new file mode 100644 index 0000000000000..0896ca890d8bf --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/BUILD @@ -0,0 +1,3 @@ +poetry_requirements( + name="poetry", +) diff --git a/llama-index-integrations/llms/llama-index-llms-yi/Makefile b/llama-index-integrations/llms/llama-index-llms-yi/Makefile new file mode 100644 index 0000000000000..b9eab05aa3706 --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/Makefile @@ -0,0 +1,17 @@ +GIT_ROOT ?= $(shell git rev-parse --show-toplevel) + +help: ## Show all Makefile targets. + @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}' + +format: ## Run code autoformatters (black). + pre-commit install + git ls-files | xargs pre-commit run black --files + +lint: ## Run linters: pre-commit (black, ruff, codespell) and mypy + pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files + +test: ## Run tests via pytest. + pytest tests + +watch-docs: ## Build and watch documentation. + sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/ diff --git a/llama-index-integrations/llms/llama-index-llms-yi/README.md b/llama-index-integrations/llms/llama-index-llms-yi/README.md new file mode 100644 index 0000000000000..ee49166191a46 --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/README.md @@ -0,0 +1 @@ +# LlamaIndex Llms Integration: Yi diff --git a/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/BUILD b/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/BUILD new file mode 100644 index 0000000000000..db46e8d6c978c --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/__init__.py b/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/__init__.py new file mode 100644 index 0000000000000..46b85a6986014 --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/__init__.py @@ -0,0 +1,4 @@ +from llama_index.llms.yi.base import Yi + + +__all__ = ["Yi"] diff --git a/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/base.py b/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/base.py new file mode 100644 index 0000000000000..694c4abe53560 --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/llama_index/llms/yi/base.py @@ -0,0 +1,198 @@ +import os +from typing import Any, Dict, Optional, Sequence, Union + +from llama_index.core.base.llms.types import ( + ChatMessage, + ChatResponse, + ChatResponseAsyncGen, + ChatResponseGen, + CompletionResponse, + CompletionResponseAsyncGen, + CompletionResponseGen, + LLMMetadata, +) +from llama_index.core.bridge.pydantic import Field +from llama_index.core.base.llms.generic_utils import ( + async_stream_completion_response_to_chat_response, + completion_response_to_chat_response, + stream_completion_response_to_chat_response, +) +from llama_index.llms.openai.base import OpenAI, Tokenizer +from transformers import AutoTokenizer + + +DEFAULT_YI_MODEL = "yi-large" +DEFAULT_YI_ENDPOINT = "https://api.01.ai/v1" + +YI_MODELS: Dict[str, int] = { + "yi-large": 16000, +} + + +def yi_modelname_to_context_size(modelname: str) -> int: + if modelname not in YI_MODELS: + raise ValueError( + f"Unknown model: {modelname}. Please provide a valid 01.AI model name." + "Known models are: " + ", ".join(YI_MODELS.keys()) + ) + + return YI_MODELS[modelname] + + +class Yi(OpenAI): + """Yi LLM. + + Examples: + `pip install llama-index-llms-yi` + + ```python + from llama_index.llms.yi import Yi + + # get api key from: https://platform.01.ai/ + llm = Yi(model="yi-large", api_key="YOUR_API_KEY") + + response = llm.chat("Hi, who are you?") + print(response) + ``` + """ + + model: str = Field(default=DEFAULT_YI_MODEL, description="The Yi model to use.") + context_window: int = Field( + default=yi_modelname_to_context_size(DEFAULT_YI_MODEL), + description=LLMMetadata.__fields__["context_window"].field_info.description, + ) + is_chat_model: bool = Field( + default=True, + description=LLMMetadata.__fields__["is_chat_model"].field_info.description, + ) + is_function_calling_model: bool = Field( + default=False, + description=LLMMetadata.__fields__[ + "is_function_calling_model" + ].field_info.description, + ) + tokenizer: Union[Tokenizer, str, None] = Field( + default=None, + description=( + "An instance of a tokenizer object that has an encode method, or the name" + " of a tokenizer model from Hugging Face. If left as None, then this" + " disables inference of max_tokens." + ), + ) + + def __init__( + self, + model: str = DEFAULT_YI_MODEL, + api_key: Optional[str] = None, + api_base: Optional[str] = DEFAULT_YI_ENDPOINT, + **kwargs: Any, + ) -> None: + api_key = api_key or os.environ.get("YI_API_KEY", None) + super().__init__( + model=model, + api_key=api_key, + api_base=api_base, + **kwargs, + ) + + @property + def metadata(self) -> LLMMetadata: + return LLMMetadata( + context_window=self.context_window, + num_output=self.max_tokens or -1, + is_chat_model=self.is_chat_model, + is_function_calling_model=self.is_function_calling_model, + model_name=self.model, + ) + + @property + def _tokenizer(self) -> Optional[Tokenizer]: + if isinstance(self.tokenizer, str): + return AutoTokenizer.from_pretrained(self.tokenizer) + return self.tokenizer + + @classmethod + def class_name(cls) -> str: + return "Yi_LLM" + + def complete( + self, prompt: str, formatted: bool = False, **kwargs: Any + ) -> CompletionResponse: + """Complete the prompt.""" + if not formatted: + prompt = self.completion_to_prompt(prompt) + + return super().complete(prompt, **kwargs) + + def stream_complete( + self, prompt: str, formatted: bool = False, **kwargs: Any + ) -> CompletionResponseGen: + """Stream complete the prompt.""" + if not formatted: + prompt = self.completion_to_prompt(prompt) + + return super().stream_complete(prompt, **kwargs) + + def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse: + """Chat with the model.""" + if not self.metadata.is_chat_model: + prompt = self.messages_to_prompt(messages) + completion_response = self.complete(prompt, formatted=True, **kwargs) + return completion_response_to_chat_response(completion_response) + + return super().chat(messages, **kwargs) + + def stream_chat( + self, messages: Sequence[ChatMessage], **kwargs: Any + ) -> ChatResponseGen: + if not self.metadata.is_chat_model: + prompt = self.messages_to_prompt(messages) + completion_response = self.stream_complete(prompt, formatted=True, **kwargs) + return stream_completion_response_to_chat_response(completion_response) + + return super().stream_chat(messages, **kwargs) + + # -- Async methods -- + + async def acomplete( + self, prompt: str, formatted: bool = False, **kwargs: Any + ) -> CompletionResponse: + """Complete the prompt.""" + if not formatted: + prompt = self.completion_to_prompt(prompt) + + return await super().acomplete(prompt, **kwargs) + + async def astream_complete( + self, prompt: str, formatted: bool = False, **kwargs: Any + ) -> CompletionResponseAsyncGen: + """Stream complete the prompt.""" + if not formatted: + prompt = self.completion_to_prompt(prompt) + + return await super().astream_complete(prompt, **kwargs) + + async def achat( + self, messages: Sequence[ChatMessage], **kwargs: Any + ) -> ChatResponse: + """Chat with the model.""" + if not self.metadata.is_chat_model: + prompt = self.messages_to_prompt(messages) + completion_response = await self.acomplete(prompt, formatted=True, **kwargs) + return completion_response_to_chat_response(completion_response) + + return await super().achat(messages, **kwargs) + + async def astream_chat( + self, messages: Sequence[ChatMessage], **kwargs: Any + ) -> ChatResponseAsyncGen: + if not self.metadata.is_chat_model: + prompt = self.messages_to_prompt(messages) + completion_response = await self.astream_complete( + prompt, formatted=True, **kwargs + ) + return async_stream_completion_response_to_chat_response( + completion_response + ) + + return await super().astream_chat(messages, **kwargs) diff --git a/llama-index-integrations/llms/llama-index-llms-yi/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-yi/pyproject.toml new file mode 100644 index 0000000000000..ea590f5d7f60d --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/pyproject.toml @@ -0,0 +1,58 @@ +[build-system] +build-backend = "poetry.core.masonry.api" +requires = ["poetry-core"] + +[tool.codespell] +check-filenames = true +check-hidden = true +# Feel free to un-skip examples, and experimental, you will just need to +# work through many typos (--write-changes and --interactive will help) +skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb" + +[tool.llamahub] +contains_example = false +import_path = "llama_index.llms.yi" + +[tool.llamahub.class_authors] +Yi = "Yimi81" + +[tool.mypy] +disallow_untyped_defs = true +# Remove venv skip when integrated with pre-commit +exclude = ["_static", "build", "examples", "notebooks", "venv"] +ignore_missing_imports = true +python_version = "3.8" + +[tool.poetry] +authors = ["Yimi81"] +description = "llama-index llms yi integration" +license = "MIT" +name = "llama-index-llms-yi" +packages = [{include = "llama_index/"}] +readme = "README.md" +version = "0.1.0" + +[tool.poetry.dependencies] +python = ">=3.8.1,<4.0" +llama-index-core = "^0.10.0" +llama-index-llms-openai = "^0.1.23" +transformers = "^4.41.2" + +[tool.poetry.group.dev.dependencies] +black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"} +codespell = {extras = ["toml"], version = ">=v2.2.6"} +ipython = "8.10.0" +jupyter = "^1.0.0" +mypy = "0.991" +pre-commit = "3.2.0" +pylint = "2.15.10" +pytest = "7.2.1" +pytest-mock = "3.11.1" +ruff = "0.0.292" +tree-sitter-languages = "^1.8.0" +types-Deprecated = ">=0.1.0" +types-PyYAML = "^6.0.12.12" +types-protobuf = "^4.24.0.4" +types-redis = "4.5.5.0" +types-requests = "2.28.11.8" # TODO: unpin when mypy>0.991 +types-setuptools = "67.1.0.0" diff --git a/llama-index-integrations/llms/llama-index-llms-yi/tests/BUILD b/llama-index-integrations/llms/llama-index-llms-yi/tests/BUILD new file mode 100644 index 0000000000000..dabf212d7e716 --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/tests/BUILD @@ -0,0 +1 @@ +python_tests() diff --git a/llama-index-integrations/llms/llama-index-llms-yi/tests/__init__.py b/llama-index-integrations/llms/llama-index-llms-yi/tests/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/llama-index-integrations/llms/llama-index-llms-yi/tests/test_llms_yi.py b/llama-index-integrations/llms/llama-index-llms-yi/tests/test_llms_yi.py new file mode 100644 index 0000000000000..dfdf41936450d --- /dev/null +++ b/llama-index-integrations/llms/llama-index-llms-yi/tests/test_llms_yi.py @@ -0,0 +1,7 @@ +from llama_index.core.base.llms.base import BaseLLM +from llama_index.llms.yi import Yi + + +def test_embedding_class(): + names_of_base_classes = [b.__name__ for b in Yi.__mro__] + assert BaseLLM.__name__ in names_of_base_classes