deepakts commited on
Commit
aef3ec1
·
1 Parent(s): f9e2643

inital version

Browse files
.gitignore ADDED
@@ -0,0 +1,164 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py,cover
50
+ .hypothesis/
51
+ .pytest_cache/
52
+ cover/
53
+
54
+ # Translations
55
+ *.mo
56
+ *.pot
57
+
58
+ # Django stuff:
59
+ *.log
60
+ local_settings.py
61
+ db.sqlite3
62
+ db.sqlite3-journal
63
+
64
+ # Flask stuff:
65
+ instance/
66
+ .webassets-cache
67
+
68
+ # Scrapy stuff:
69
+ .scrapy
70
+
71
+ # Sphinx documentation
72
+ docs/_build/
73
+
74
+ # PyBuilder
75
+ .pybuilder/
76
+ target/
77
+
78
+ # Jupyter Notebook
79
+ .ipynb_checkpoints
80
+
81
+ # IPython
82
+ profile_default/
83
+ ipython_config.py
84
+
85
+ # pyenv
86
+ # For a library or package, you might want to ignore these files since the code is
87
+ # intended to run in multiple environments; otherwise, check them in:
88
+ # .python-version
89
+
90
+ # pipenv
91
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
93
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
94
+ # install all needed dependencies.
95
+ #Pipfile.lock
96
+
97
+ # poetry
98
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
99
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
100
+ # commonly ignored for libraries.
101
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
102
+ #poetry.lock
103
+
104
+ # pdm
105
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
106
+ #pdm.lock
107
+ # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
108
+ # in version control.
109
+ # https://pdm.fming.dev/#use-with-ide
110
+ .pdm.toml
111
+
112
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
113
+ __pypackages__/
114
+
115
+ # Celery stuff
116
+ celerybeat-schedule
117
+ celerybeat.pid
118
+
119
+ # SageMath parsed files
120
+ *.sage.py
121
+
122
+ # Environments
123
+ .env
124
+ .venv
125
+ env/
126
+ venv/
127
+ ENV/
128
+ env.bak/
129
+ venv.bak/
130
+ .chainlit
131
+ .chainlit/
132
+ chainlit/
133
+ wandb/
134
+ wandb
135
+ # Spyder project settings
136
+ .spyderproject
137
+ .spyproject
138
+
139
+ # Rope project settings
140
+ .ropeproject
141
+
142
+ # mkdocs documentation
143
+ /site
144
+
145
+ # mypy
146
+ .mypy_cache/
147
+ .dmypy.json
148
+ dmypy.json
149
+
150
+ # Pyre type checker
151
+ .pyre/
152
+
153
+ # pytype static type analyzer
154
+ .pytype/
155
+
156
+ # Cython debug symbols
157
+ cython_debug/
158
+
159
+ # PyCharm
160
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
161
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
162
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
163
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
164
+ #.idea/
Dockerfile ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.9
2
+ RUN useradd -m -u 1000 user
3
+ USER user
4
+ ENV HOME=/home/user \
5
+ PATH=/home/user/.local/bin:$PATH
6
+ WORKDIR $HOME/app
7
+ COPY --chown=user . $HOME/app
8
+ COPY ./requirements.txt ~/app/requirements.txt
9
+ RUN pip install -r requirements.txt
10
+ COPY . .
11
+ CMD ["chainlit", "run", "app.py", "--port", "7860"]
app.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import openai
3
+ import platform
4
+ import wandb
5
+ import chainlit as cl #importing chainlit for our app
6
+ from chainlit.input_widget import Select, Switch, Slider #importing chainlit settings selection tools
7
+ from chainlit.prompt import Prompt, PromptMessage #importing prompt tools
8
+ from chainlit.playground.providers import ChatOpenAI #importing ChatOpenAI tools
9
+ import asyncio
10
+ from makersutil.text_utils import TextFileLoader, CharacterTextSplitter
11
+ from makersutil.vectordatabase import VectorDatabase
12
+ from makersutil.retrievalAugmentedQAPipeline import RetrievalAugmentedQAPipeline
13
+ from makersutil.openai_utils.chatmodel import ChatOpenAI
14
+
15
+
16
+
17
+ @cl.on_chat_start # marks a function that will be executed at the start of a user session
18
+ async def start_chat():
19
+ pass
20
+ # nothing for now
21
+ # settings = {
22
+ # "model": "gpt-3.5-turbo",
23
+ # "temperature": 0,
24
+ # "max_tokens": 500,
25
+ # "top_p": 1,
26
+ # "frequency_penalty": 0,
27
+ # "presence_penalty": 0,
28
+ # }
29
+
30
+ @cl.on_message # marks a function that should be run each time the chatbot receives a message from a user
31
+ async def main(message: str):
32
+ wandb.init(project="KingLearbook")
33
+ msg = cl.Message(content="")
34
+ text_loader = TextFileLoader("data/KingLear.txt")
35
+ documents = text_loader.load_documents()
36
+ text_splitter = CharacterTextSplitter()
37
+ split_documents = text_splitter.split_texts(documents)
38
+ vector_db = VectorDatabase()
39
+ asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
40
+ vector_db = asyncio.run(vector_db.abuild_from_list(split_documents))
41
+
42
+ chat_openai = ChatOpenAI()
43
+
44
+ retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
45
+ vector_db_retriever=vector_db,
46
+ llm=chat_openai,
47
+ wandb_project="KingLearbook",
48
+ )
49
+
50
+ msg.content = retrieval_augmented_qa_pipeline.run_pipeline(message)
51
+ await msg.send()
chainlit.md ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ # Welcome to Chainlit! 🚀🤖
2
+
3
+ Hi there, do you want to know about King Lear , lets Chat !!
data/KingLear.txt ADDED
The diff for this file is too large to render. See raw diff
 
makersutil/__init__.py ADDED
File without changes
makersutil/openai_utils/__init__.py ADDED
File without changes
makersutil/openai_utils/chatmodel.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import openai
2
+ from dotenv import load_dotenv
3
+ import os
4
+
5
+ load_dotenv()
6
+
7
+
8
+ class ChatOpenAI:
9
+ def __init__(self, model_name: str = "gpt-3.5-turbo"):
10
+ self.model_name = model_name
11
+ self.openai_api_key = os.getenv("OPENAI_API_KEY")
12
+ if self.openai_api_key is None:
13
+ raise ValueError("OPENAI_API_KEY is not set")
14
+
15
+ def run(self, messages, text_only: bool = True):
16
+ if not isinstance(messages, list):
17
+ raise ValueError("messages must be a list")
18
+
19
+ openai.api_key = self.openai_api_key
20
+ openai.temperature = 0
21
+ openai.max_tokens = 500
22
+ openai.top_p = 1
23
+ openai.frequency_penalty = 0
24
+ openai.presence_penalty = 0
25
+ response = openai.ChatCompletion.create(
26
+ model=self.model_name, messages=messages
27
+ )
28
+
29
+ if text_only:
30
+ return response.choices[0].message.content
31
+
32
+ return response
makersutil/openai_utils/embedding.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from dotenv import load_dotenv
2
+ from openai.embeddings_utils import (
3
+ get_embeddings,
4
+ aget_embeddings,
5
+ get_embedding,
6
+ aget_embedding,
7
+ )
8
+ import openai
9
+ from typing import List
10
+ import os
11
+ import asyncio
12
+
13
+
14
+ class EmbeddingModel:
15
+ def __init__(self, embeddings_model_name: str = "text-embedding-ada-002"):
16
+ load_dotenv()
17
+ self.openai_api_key = os.getenv("OPENAI_API_KEY")
18
+
19
+ if self.openai_api_key is None:
20
+ raise ValueError(
21
+ "OPENAI_API_KEY environment variable is not set. Please set it to your OpenAI API key."
22
+ )
23
+ openai.api_key = self.openai_api_key
24
+ self.embeddings_model_name = embeddings_model_name
25
+
26
+ async def async_get_embeddings(self, list_of_text: List[str]) -> List[List[float]]:
27
+ return await aget_embeddings(
28
+ list_of_text=list_of_text, engine=self.embeddings_model_name
29
+ )
30
+
31
+ async def async_get_embedding(self, text: str) -> List[float]:
32
+ return await aget_embedding(text=text, engine=self.embeddings_model_name)
33
+
34
+ def get_embeddings(self, list_of_text: List[str]) -> List[List[float]]:
35
+ return get_embeddings(
36
+ list_of_text=list_of_text, engine=self.embeddings_model_name
37
+ )
38
+
39
+ def get_embedding(self, text: str) -> List[float]:
40
+ return get_embedding(text=text, engine=self.embeddings_model_name)
41
+
42
+
43
+ if __name__ == "__main__":
44
+ embedding_model = EmbeddingModel()
45
+ print(embedding_model.get_embedding("Hello, world!"))
46
+ print(embedding_model.get_embeddings(["Hello, world!", "Goodbye, world!"]))
47
+ print(asyncio.run(embedding_model.async_get_embedding("Hello, world!")))
48
+ print(
49
+ asyncio.run(
50
+ embedding_model.async_get_embeddings(["Hello, world!", "Goodbye, world!"])
51
+ )
52
+ )
makersutil/openai_utils/prompts.py ADDED
@@ -0,0 +1,75 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+
3
+
4
+ class BasePrompt:
5
+ def __init__(self, prompt):
6
+ """
7
+ Initializes the BasePrompt object with a prompt template.
8
+
9
+ :param prompt: A string that can contain placeholders within curly braces
10
+ """
11
+ self.prompt = prompt
12
+ self._pattern = re.compile(r"\{([^}]+)\}")
13
+
14
+ def format_prompt(self, **kwargs):
15
+ """
16
+ Formats the prompt string using the keyword arguments provided.
17
+
18
+ :param kwargs: The values to substitute into the prompt string
19
+ :return: The formatted prompt string
20
+ """
21
+ matches = self._pattern.findall(self.prompt)
22
+ return self.prompt.format(**{match: kwargs.get(match, "") for match in matches})
23
+
24
+ def get_input_variables(self):
25
+ """
26
+ Gets the list of input variable names from the prompt string.
27
+
28
+ :return: List of input variable names
29
+ """
30
+ return self._pattern.findall(self.prompt)
31
+
32
+
33
+ class RolePrompt(BasePrompt):
34
+ def __init__(self, prompt, role: str):
35
+ """
36
+ Initializes the RolePrompt object with a prompt template and a role.
37
+
38
+ :param prompt: A string that can contain placeholders within curly braces
39
+ :param role: The role for the message ('system', 'user', or 'assistant')
40
+ """
41
+ super().__init__(prompt)
42
+ self.role = role
43
+
44
+ def create_message(self, **kwargs):
45
+ """
46
+ Creates a message dictionary with a role and a formatted message.
47
+
48
+ :param kwargs: The values to substitute into the prompt string
49
+ :return: Dictionary containing the role and the formatted message
50
+ """
51
+ return {"role": self.role, "content": self.format_prompt(**kwargs)}
52
+
53
+
54
+ class SystemRolePrompt(RolePrompt):
55
+ def __init__(self, prompt: str):
56
+ super().__init__(prompt, "system")
57
+
58
+
59
+ class UserRolePrompt(RolePrompt):
60
+ def __init__(self, prompt: str):
61
+ super().__init__(prompt, "user")
62
+
63
+
64
+ class AssistantRolePrompt(RolePrompt):
65
+ def __init__(self, prompt: str):
66
+ super().__init__(prompt, "assistant")
67
+
68
+
69
+ if __name__ == "__main__":
70
+ prompt = BasePrompt("Hello {name}, you are {age} years old")
71
+ print(prompt.format_prompt(name="John", age=30))
72
+
73
+ prompt = SystemRolePrompt("Hello {name}, you are {age} years old")
74
+ print(prompt.create_message(name="John", age=30))
75
+ print(prompt.get_input_variables())
makersutil/retrievalAugmentedQAPipeline.py ADDED
@@ -0,0 +1,88 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from makersutil.openai_utils.prompts import (
3
+ UserRolePrompt,
4
+ SystemRolePrompt,
5
+ AssistantRolePrompt,
6
+ )
7
+
8
+ from makersutil.vectordatabase import VectorDatabase
9
+ from makersutil.openai_utils.chatmodel import ChatOpenAI
10
+ import datetime
11
+ from wandb.sdk.data_types.trace_tree import Trace
12
+
13
+
14
+ RAQA_PROMPT_TEMPLATE = """
15
+ Use the provided context to answer the user's query.
16
+
17
+ You may not answer the user's query unless there is specific context in the following text.
18
+
19
+ If you do not know the answer, or cannot answer, please respond with "I don't know".
20
+
21
+ Context:
22
+ {context}
23
+ """
24
+
25
+ raqa_prompt = SystemRolePrompt(RAQA_PROMPT_TEMPLATE)
26
+
27
+ USER_PROMPT_TEMPLATE = """
28
+ User Query:
29
+ {user_query}
30
+ """
31
+
32
+ user_prompt = UserRolePrompt(USER_PROMPT_TEMPLATE)
33
+
34
+ class RetrievalAugmentedQAPipeline:
35
+ def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase, wandb_project = None) -> None:
36
+ self.llm = llm
37
+ self.vector_db_retriever = vector_db_retriever
38
+ self.wandb_project = wandb_project
39
+
40
+ def run_pipeline(self, user_query: str) -> str:
41
+ context_list = self.vector_db_retriever.search_by_text(user_query, k=4)
42
+
43
+ context_prompt = ""
44
+ for context in context_list:
45
+ context_prompt += context[0] + "\n"
46
+
47
+ formatted_system_prompt = raqa_prompt.create_message(context=context_prompt)
48
+
49
+ formatted_user_prompt = user_prompt.create_message(user_query=user_query)
50
+
51
+ start_time = datetime.datetime.now().timestamp() * 1000
52
+
53
+ try:
54
+ openai_response = self.llm.run([formatted_system_prompt, formatted_user_prompt], text_only=False)
55
+ end_time = datetime.datetime.now().timestamp() * 1000
56
+ status = "success"
57
+ status_message = (None, )
58
+ response_text = openai_response.choices[0].message.content
59
+ token_usage = openai_response["usage"].to_dict()
60
+ model = openai_response["model"]
61
+
62
+ except Exception as e:
63
+ end_time = datetime.datetime.now().timestamp() * 1000
64
+ status = "error"
65
+ status_message = str(e)
66
+ response_text = ""
67
+ token_usage = {}
68
+ model = ""
69
+
70
+ if self.wandb_project:
71
+ root_span = Trace(
72
+ name="root_span",
73
+ kind="llm",
74
+ status_code=status,
75
+ status_message=status_message,
76
+ start_time_ms=start_time,
77
+ end_time_ms=end_time,
78
+ metadata={
79
+ "token_usage" : token_usage,
80
+ "model_name" : model
81
+ },
82
+ inputs= {"system_prompt" : formatted_system_prompt, "user_prompt" : formatted_user_prompt},
83
+ outputs= {"response" : response_text}
84
+ )
85
+
86
+ root_span.log(name="openai_trace")
87
+
88
+ return response_text if response_text else "We ran into an error. Please try again later. Full Error Message: " + status_message
makersutil/text_utils.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from typing import List
3
+
4
+
5
+ class TextFileLoader:
6
+ def __init__(self, path: str, encoding: str = "utf-8"):
7
+ self.documents = []
8
+ self.path = path
9
+ self.encoding = encoding
10
+
11
+ def load(self):
12
+ if os.path.isdir(self.path):
13
+ self.load_directory()
14
+ elif os.path.isfile(self.path) and self.path.endswith(".txt"):
15
+ self.load_file()
16
+ else:
17
+ raise ValueError(
18
+ "Provided path is neither a valid directory nor a .txt file."
19
+ )
20
+
21
+ def load_file(self):
22
+ with open(self.path, "r", encoding=self.encoding) as f:
23
+ self.documents.append(f.read())
24
+
25
+ def load_directory(self):
26
+ for root, _, files in os.walk(self.path):
27
+ for file in files:
28
+ if file.endswith(".txt"):
29
+ with open(
30
+ os.path.join(root, file), "r", encoding=self.encoding
31
+ ) as f:
32
+ self.documents.append(f.read())
33
+
34
+ def load_documents(self):
35
+ self.load()
36
+ return self.documents
37
+
38
+
39
+ class CharacterTextSplitter:
40
+ def __init__(
41
+ self,
42
+ chunk_size: int = 1000,
43
+ chunk_overlap: int = 200,
44
+ ):
45
+ assert (
46
+ chunk_size > chunk_overlap
47
+ ), "Chunk size must be greater than chunk overlap"
48
+
49
+ self.chunk_size = chunk_size
50
+ self.chunk_overlap = chunk_overlap
51
+
52
+ def split(self, text: str) -> List[str]:
53
+ chunks = []
54
+ for i in range(0, len(text), self.chunk_size - self.chunk_overlap):
55
+ chunks.append(text[i : i + self.chunk_size])
56
+ return chunks
57
+
58
+ def split_texts(self, texts: List[str]) -> List[str]:
59
+ chunks = []
60
+ for text in texts:
61
+ chunks.extend(self.split(text))
62
+ return chunks
63
+
64
+
65
+ if __name__ == "__main__":
66
+ loader = TextFileLoader("data/KingLear.txt")
67
+ loader.load()
68
+ splitter = CharacterTextSplitter()
69
+ chunks = splitter.split_texts(loader.documents)
70
+ print(len(chunks))
71
+ print(chunks[0])
72
+ print("--------")
73
+ print(chunks[1])
74
+ print("--------")
75
+ print(chunks[-2])
76
+ print("--------")
77
+ print(chunks[-1])
makersutil/vectordatabase.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ from collections import defaultdict
3
+ from typing import List, Tuple, Callable
4
+ from makersutil.openai_utils.embedding import EmbeddingModel
5
+ import asyncio
6
+
7
+
8
+ def cosine_similarity(vector_a: np.array, vector_b: np.array) -> float:
9
+ """Computes the cosine similarity between two vectors."""
10
+ dot_product = np.dot(vector_a, vector_b)
11
+ norm_a = np.linalg.norm(vector_a)
12
+ norm_b = np.linalg.norm(vector_b)
13
+ return dot_product / (norm_a * norm_b)
14
+
15
+
16
+ class VectorDatabase:
17
+ def __init__(self, embedding_model: EmbeddingModel = None):
18
+ self.vectors = defaultdict(np.array)
19
+ self.embedding_model = embedding_model or EmbeddingModel()
20
+
21
+ def insert(self, key: str, vector: np.array) -> None:
22
+ self.vectors[key] = vector
23
+
24
+ def search(
25
+ self,
26
+ query_vector: np.array,
27
+ k: int,
28
+ distance_measure: Callable = cosine_similarity,
29
+ ) -> List[Tuple[str, float]]:
30
+ scores = [
31
+ (key, distance_measure(query_vector, vector))
32
+ for key, vector in self.vectors.items()
33
+ ]
34
+ return sorted(scores, key=lambda x: x[1], reverse=True)[:k]
35
+
36
+ def search_by_text(
37
+ self,
38
+ query_text: str,
39
+ k: int,
40
+ distance_measure: Callable = cosine_similarity,
41
+ return_as_text: bool = False,
42
+ ) -> List[Tuple[str, float]]:
43
+ query_vector = self.embedding_model.get_embedding(query_text)
44
+ results = self.search(query_vector, k, distance_measure)
45
+ return [result[0] for result in results] if return_as_text else results
46
+
47
+ def retrieve_from_key(self, key: str) -> np.array:
48
+ return self.vectors.get(key, None)
49
+
50
+ async def abuild_from_list(self, list_of_text: List[str]) -> "VectorDatabase":
51
+ embeddings = await self.embedding_model.async_get_embeddings(list_of_text)
52
+ for text, embedding in zip(list_of_text, embeddings):
53
+ self.insert(text, np.array(embedding))
54
+ return self
55
+
56
+
57
+ if __name__ == "__main__":
58
+ list_of_text = [
59
+ "I like to eat broccoli and bananas.",
60
+ "I ate a banana and spinach smoothie for breakfast.",
61
+ "Chinchillas and kittens are cute.",
62
+ "My sister adopted a kitten yesterday.",
63
+ "Look at this cute hamster munching on a piece of broccoli.",
64
+ ]
65
+
66
+ vector_db = VectorDatabase()
67
+ vector_db = asyncio.run(vector_db.abuild_from_list(list_of_text))
68
+ k = 2
69
+
70
+ searched_vector = vector_db.search_by_text("I think fruit is awesome!", k=k)
71
+ print(f"Closest {k} vector(s):", searched_vector)
72
+
73
+ retrieved_vector = vector_db.retrieve_from_key(
74
+ "I like to eat broccoli and bananas."
75
+ )
76
+ print("Retrieved vector:", retrieved_vector)
77
+
78
+ relevant_texts = vector_db.search_by_text(
79
+ "I think fruit is awesome!", k=k, return_as_text=True
80
+ )
81
+ print(f"Closest {k} text(s):", relevant_texts)
requirements.txt ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ numpy==1.25.2
2
+ openai==0.27.8
3
+ python-dotenv==1.0.0
4
+ pandas
5
+ scikit-learn
6
+ ipykernel
7
+ matplotlib
8
+ plotly
9
+ chainlit
10
+ wandb