Public/kotaemon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Huggingface embeddings and Cohere embeddings (#63)

Browse Source 
* Add huggingface embeddings and cohere embeddings
* Update openai interface and the mock for newer OpenAI SDK

---------

Co-authored-by: trducng <trungduc1992@gmail.com>
This commit is contained in:
ian_Cin 2023-11-10 09:38:30 +07:00 committed by GitHub
parent 9035e25666
commit 6095526dc7
13 changed files with 249 additions and 138 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
knowledgehub/embeddings/cohere.py Normal file
View File

@ -0,0 +1,12 @@
from langchain.embeddings import CohereEmbeddings as LCCohereEmbeddings
from kotaemon.embeddings.base import LangchainEmbeddings
class CohereEmbdeddings(LangchainEmbeddings):
"""Cohere embeddings.
This class wraps around the Langchain CohereEmbeddings class.
"""
_lc_class = LCCohereEmbeddings

12
knowledgehub/embeddings/huggingface.py Normal file
View File

@ -0,0 +1,12 @@
from langchain.embeddings import HuggingFaceBgeEmbeddings as LCHuggingFaceEmbeddings
from kotaemon.embeddings.base import LangchainEmbeddings
class HuggingFaceEmbeddings(LangchainEmbeddings):
"""HuggingFace embeddings
This class wraps around the Langchain HuggingFaceEmbeddings class
"""
_lc_class = LCHuggingFaceEmbeddings

4
knowledgehub/loaders/base.py
View File

@ -8,7 +8,7 @@ from ..base import BaseComponent
from ..documents.base import Document
class AutoReader(BaseComponent, BaseReader):
class AutoReader(BaseComponent):
"""General auto reader for a variety of files. (based on llama-hub)"""
def __init__(self, reader_type: Union[str, Type[BaseReader]]) -> None:
@ -31,7 +31,7 @@ class AutoReader(BaseComponent, BaseReader):
return self.load_data(file=file, **kwargs)
class LIBaseReader(BaseComponent, BaseReader):
class LIBaseReader(BaseComponent):
_reader_class: Type[BaseReader]
def __init__(self, *args, **kwargs):

4
knowledgehub/pipelines/agents/rewoo/planner.py
View File

@ -75,8 +75,8 @@ class Planner(BaseComponent):
try:
response = self.model(prompt)
output.info("Planner run successful.")
except ValueError:
except ValueError as e:
output.error("Planner failed to retrieve response from LLM")
raise ValueError("Planner failed to retrieve response from LLM")
raise ValueError("Planner failed to retrieve response from LLM") from e
return response

3
setup.py
View File

@ -59,6 +59,9 @@ setuptools.setup(
"pytest-mock",
"unstructured[pdf]",
"farm-haystack==1.19.0",
"sentence_transformers",
"cohere",
"pypdf",
],
},
entry_points={"console_scripts": ["kh=kotaemon.cli:main"]},

129
tests/test_agent.py
View File

@ -1,6 +1,7 @@
from unittest.mock import patch
import pytest
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.chats.openai import AzureChatOpenAI
from kotaemon.pipelines.agents.react import ReactAgent
@ -14,24 +15,30 @@ from kotaemon.pipelines.tools import (
FINAL_RESPONSE_TEXT = "Hello Cinnamon AI!"
_openai_chat_completion_responses_rewoo = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
ChatCompletion.parse_obj(
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in [
(
"#Plan1: Search for Cinnamon AI company on Google\n"
@ -44,23 +51,28 @@ _openai_chat_completion_responses_rewoo = [
]
_openai_chat_completion_responses_react = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
ChatCompletion.parse_obj(
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in [
(
"I don't have prior knowledge about Cinnamon AI company, "
@ -82,23 +94,28 @@ _openai_chat_completion_responses_react = [
]
_openai_chat_completion_responses_react_langchain_tool = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
ChatCompletion.parse_obj(
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in [
(
"I don't have prior knowledge about Cinnamon AI company, "
@ -123,7 +140,7 @@ _openai_chat_completion_responses_react_langchain_tool = [
@pytest.fixture
def llm():
return AzureChatOpenAI(
openai_api_base="https://dummy.openai.azure.com/",
azure_endpoint="https://dummy.openai.azure.com/",
openai_api_key="dummy",
openai_api_version="2023-03-15-preview",
deployment_name="dummy-q2",
@ -132,7 +149,7 @@ def llm():
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_rewoo,
)
def test_rewoo_agent(openai_completion, llm, mock_google_search):
@ -150,7 +167,7 @@ def test_rewoo_agent(openai_completion, llm, mock_google_search):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_react,
)
def test_react_agent(openai_completion, llm, mock_google_search):
@ -167,7 +184,7 @@ def test_react_agent(openai_completion, llm, mock_google_search):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_react,
)
def test_react_agent_langchain(openai_completion, llm, mock_google_search):
@ -191,7 +208,7 @@ def test_react_agent_langchain(openai_completion, llm, mock_google_search):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_responses_react_langchain_tool,
)
def test_react_agent_with_langchain_tools(openai_completion, llm):

47
tests/test_cot.py
View File

@ -1,32 +1,39 @@
from unittest.mock import patch
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.chats.openai import AzureChatOpenAI
from kotaemon.pipelines.cot import ManualSequentialChainOfThought, Thought
_openai_chat_completion_response = [
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
ChatCompletion.parse_obj(
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": text,
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
for text in ["Bonjour", "こんにちは (Konnichiwa)"]
]
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_response,
)
def test_cot_plus_operator(openai_completion):
@ -58,7 +65,7 @@ def test_cot_plus_operator(openai_completion):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_response,
)
def test_cot_manual(openai_completion):
@ -88,7 +95,7 @@ def test_cot_manual(openai_completion):
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=_openai_chat_completion_response,
)
def test_cot_with_termination_callback(openai_completion):

45
tests/test_embedding_models.py
View File

@ -2,6 +2,8 @@ import json
from pathlib import Path
from unittest.mock import patch
from kotaemon.embeddings.cohere import CohereEmbdeddings
from kotaemon.embeddings.huggingface import HuggingFaceEmbeddings
from kotaemon.embeddings.openai import AzureOpenAIEmbeddings
with open(Path(__file__).parent / "resources" / "embedding_openai_batch.json") as f:
@ -12,7 +14,7 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
@patch(
"openai.api_resources.embedding.Embedding.create",
"openai.resources.embeddings.Embeddings.create",
side_effect=lambda *args, **kwargs: openai_embedding,
)
def test_azureopenai_embeddings_raw(openai_embedding_call):
@ -29,7 +31,7 @@ def test_azureopenai_embeddings_raw(openai_embedding_call):
@patch(
"openai.api_resources.embedding.Embedding.create",
"openai.resources.embeddings.Embeddings.create",
side_effect=lambda *args, **kwargs: openai_embedding_batch,
)
def test_azureopenai_embeddings_batch_raw(openai_embedding_call):
@ -44,3 +46,42 @@ def test_azureopenai_embeddings_batch_raw(openai_embedding_call):
assert isinstance(output[0], list)
assert isinstance(output[0][0], float)
openai_embedding_call.assert_called()
@patch(
"sentence_transformers.SentenceTransformer",
side_effect=lambda *args, **kwargs: None,
)
@patch(
"langchain.embeddings.huggingface.HuggingFaceBgeEmbeddings.embed_query",
side_effect=lambda *args, **kwargs: [1.0, 2.1, 3.2],
)
def test_huggingface_embddings(
langchain_huggingface_embedding_call, sentence_transformers_init
):
model = HuggingFaceEmbeddings(
model_name="intfloat/multilingual-e5-large",
model_kwargs={"device": "cpu"},
encode_kwargs={"normalize_embeddings": False},
)
output = model("Hello World")
assert isinstance(output, list)
assert isinstance(output[0], float)
sentence_transformers_init.assert_called()
langchain_huggingface_embedding_call.assert_called()
@patch(
"langchain.embeddings.cohere.CohereEmbeddings.embed_query",
side_effect=lambda *args, **kwargs: [1.0, 2.1, 3.2],
)
def test_cohere_embddings(langchain_cohere_embedding_call):
model = CohereEmbdeddings(
model="embed-english-light-v2.0", cohere_api_key="my-api-key"
)
output = model("Hello World")
assert isinstance(output, list)
assert isinstance(output[0], float)
langchain_cohere_embedding_call.assert_called()

4
tests/test_indexing_retrieval.py
View File

@ -3,7 +3,7 @@ from pathlib import Path
from typing import cast
import pytest
from openai.api_resources.embedding import Embedding
from openai.resources.embeddings import Embeddings
from kotaemon.docstores import InMemoryDocumentStore
from kotaemon.documents.base import Document
@ -18,7 +18,7 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
@pytest.fixture(scope="function")
def mock_openai_embedding(monkeypatch):
monkeypatch.setattr(Embedding, "create", lambda *args, **kwargs: openai_embedding)
monkeypatch.setattr(Embeddings, "create", lambda *args, **kwargs: openai_embedding)
def test_indexing(mock_openai_embedding, tmp_path):

43
tests/test_llms_chat_models.py
View File

@ -2,31 +2,37 @@ from unittest.mock import patch
from langchain.chat_models import AzureChatOpenAI as AzureChatOpenAILC
from langchain.schema.messages import AIMessage, HumanMessage, SystemMessage
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.base import LLMInterface
from kotaemon.llms.chats.openai import AzureChatOpenAI
_openai_chat_completion_response = {
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "Hello! How can I assist you today?",
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
_openai_chat_completion_response = ChatCompletion.parse_obj(
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "Hello! How can I assist you today?",
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_chat_completion_response,
)
def test_azureopenai_model(openai_completion):
@ -36,7 +42,6 @@ def test_azureopenai_model(openai_completion):
openai_api_version="2023-03-15-preview",
deployment_name="gpt35turbo",
temperature=0,
request_timeout=60,
)
assert isinstance(
model.agent, AzureChatOpenAILC

33
tests/test_llms_completion_models.py
View File

@ -2,24 +2,33 @@ from unittest.mock import patch
from langchain.llms import AzureOpenAI as AzureOpenAILC
from langchain.llms import OpenAI as OpenAILC
from openai.types.completion import Completion
from kotaemon.llms.base import LLMInterface
from kotaemon.llms.completions.openai import AzureOpenAI, OpenAI
_openai_completion_response = {
"id": "cmpl-7qyNoIo6gRSCJR0hi8o3ZKBH4RkJ0",
"object": "sample text_completion",
"created": 1392751226,
"model": "gpt-35-turbo",
"choices": [
{"text": "completion", "index": 0, "finish_reason": "length", "logprobs": None}
],
"usage": {"completion_tokens": 20, "prompt_tokens": 2, "total_tokens": 22},
}
_openai_completion_response = Completion.parse_obj(
{
"id": "cmpl-7qyNoIo6gRSCJR0hi8o3ZKBH4RkJ0",
"object": "text_completion",
"created": 1392751226,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"text": "completion",
"index": 0,
"finish_reason": "length",
"logprobs": None,
}
],
"usage": {"completion_tokens": 20, "prompt_tokens": 2, "total_tokens": 22},
}
)
@patch(
"openai.api_resources.completion.Completion.create",
"openai.resources.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_completion_response,
)
def test_azureopenai_model(openai_completion):
@ -47,7 +56,7 @@ def test_azureopenai_model(openai_completion):
@patch(
"openai.api_resources.completion.Completion.create",
"openai.resources.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_completion_response,
)
def test_openai_model(openai_completion):

47
tests/test_qa.py
View File

@ -3,7 +3,8 @@ from pathlib import Path
from unittest.mock import patch
import pytest
from openai.api_resources.embedding import Embedding
from openai.resources.embeddings import Embeddings
from openai.types.chat.chat_completion import ChatCompletion
from kotaemon.llms.chats.openai import AzureChatOpenAI
from kotaemon.pipelines.ingest import ReaderIndexingPipeline
@ -12,32 +13,37 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
openai_embedding = json.load(f)
_openai_chat_completion_response = {
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "Hello! How can I assist you today?",
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
_openai_chat_completion_response = ChatCompletion.parse_obj(
{
"id": "chatcmpl-7qyuw6Q1CFCpcKsMdFkmUPUa7JP2x",
"object": "chat.completion",
"created": 1692338378,
"model": "gpt-35-turbo",
"system_fingerprint": None,
"choices": [
{
"index": 0,
"finish_reason": "stop",
"message": {
"role": "assistant",
"content": "Hello! How can I assist you today?",
"function_call": None,
"tool_calls": None,
},
}
],
"usage": {"completion_tokens": 9, "prompt_tokens": 10, "total_tokens": 19},
}
)
@pytest.fixture(scope="function")
def mock_openai_embedding(monkeypatch):
monkeypatch.setattr(Embedding, "create", lambda *args, **kwargs: openai_embedding)
monkeypatch.setattr(Embeddings, "create", lambda *args, **kwargs: openai_embedding)
@patch(
"openai.api_resources.chat_completion.ChatCompletion.create",
"openai.resources.chat.completions.Completions.create",
side_effect=lambda *args, **kwargs: _openai_chat_completion_response,
)
def test_ingest_pipeline(patch, mock_openai_embedding, tmp_path):
@ -61,7 +67,6 @@ def test_ingest_pipeline(patch, mock_openai_embedding, tmp_path):
openai_api_version="2023-03-15-preview",
deployment_name="gpt35turbo",
temperature=0,
request_timeout=60,
)
qa_pipeline = indexing_pipeline.to_qa_pipeline(llm=llm, openai_api_key="some-key")
response = qa_pipeline("Summarize this document.")

4
tests/test_tools.py
View File

@ -2,7 +2,7 @@ import json
from pathlib import Path
import pytest
from openai.api_resources.embedding import Embedding
from openai.resources.embeddings import Embeddings
from kotaemon.docstores import InMemoryDocumentStore
from kotaemon.documents.base import Document
@ -18,7 +18,7 @@ with open(Path(__file__).parent / "resources" / "embedding_openai.json") as f:
@pytest.fixture(scope="function")
def mock_openai_embedding(monkeypatch):
monkeypatch.setattr(Embedding, "create", lambda *args, **kwargs: openai_embedding)
monkeypatch.setattr(Embeddings, "create", lambda *args, **kwargs: openai_embedding)
def test_google_tool(mock_google_search):