teddyllm commited on
Commit
bd3532f
·
verified ·
1 Parent(s): 06217a7

Upload 20 files

Browse files
.chromadb/20cb0d5e-499b-4d31-8a46-9a4bf83219b0/data_level0.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a13e72541800c513c73dccea69f79e39cf4baef4fa23f7e117c0d6b0f5f99670
3
+ size 3212000
.chromadb/20cb0d5e-499b-4d31-8a46-9a4bf83219b0/header.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0ec6df10978b056a10062ed99efeef2702fa4a1301fad702b53dd2517103c746
3
+ size 100
.chromadb/20cb0d5e-499b-4d31-8a46-9a4bf83219b0/length.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:939f9b658043836d45463f2f0a74346c2d2431220230ae6d8c2ad709814800bb
3
+ size 4000
.chromadb/20cb0d5e-499b-4d31-8a46-9a4bf83219b0/link_lists.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
3
+ size 0
.chromadb/chroma.sqlite3 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3484a9a12bf3911aa0a9d714f2a7b1c5a6b683276c2146cd4342065c375d29d5
3
+ size 1622016
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ .chromadb/chroma.sqlite3 filter=lfs diff=lfs merge=lfs -text
SmartSearch/README.md ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Smart Search
2
+ ### Enhancing search experience with natural language
3
+
4
+
5
+ The search should be able to support multiple llms providers
6
+
7
+ * create embeddings
8
+ * can do hybrid search
9
+ * semantic search
10
+ * keyword search
11
+
12
+
13
+
14
+
15
+
16
+
17
+
18
+
19
+
SmartSearch/database/__init__.py ADDED
File without changes
SmartSearch/database/annoydb.py ADDED
@@ -0,0 +1,89 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Dict, Union, Any
2
+
3
+ import numpy as np
4
+ from annoy import AnnoyIndex
5
+
6
+ from .vector_store import VectorStore
7
+
8
+ class AnnoyDB(VectorStore):
9
+ def __init__(
10
+ self,
11
+ embedding_dim: int,
12
+ metric: str = 'angular'
13
+ ) -> None:
14
+ self.documents = []
15
+ self.metadata = []
16
+ self.embedding_dim = embedding_dim
17
+
18
+ self.index = AnnoyIndex(embedding_dim, metric)
19
+ self.index_built = False
20
+
21
+ def add_document(self, text: str, metadata: Dict[str, Any] = None):
22
+ """
23
+ Add a document to the search index.
24
+
25
+ Args:
26
+ text: The document text
27
+ metadata: Optional metadata about the document
28
+ """
29
+ self.documents.append(text)
30
+ self.metadata.append(metadata or {})
31
+
32
+ # Generate embedding using Sentence Transformers
33
+ embedding = self.model.encode(text, show_progress_bar=False)
34
+
35
+ # Add to Annoy index
36
+ index_id = len(self.documents) - 1
37
+ self.index.add_item(index_id, embedding)
38
+ self.index_built = False
39
+
40
+ def add_documents(self, texts: List[str], embeddings: np.array, metadata_list: List[Dict[str, Any]] = None):
41
+ """
42
+ Batch add documents to the search index.
43
+
44
+ Args:
45
+ texts: List of document texts
46
+ metadata_list: Optional list of metadata dictionaries
47
+ """
48
+ if metadata_list is None:
49
+ metadata_list = [{} for _ in texts]
50
+
51
+ # Add documents and embeddings
52
+ print("Adding to index...")
53
+ for i, (text, metadata, embedding) in enumerate(zip(texts, metadata_list, embeddings)):
54
+ self.documents.append(text)
55
+ self.metadata.append(metadata)
56
+ self.index.add_item(len(self.documents) - 1, embedding)
57
+
58
+ self.index_built = False
59
+ print("Done")
60
+
61
+ def add_data(self, embedding: np.ndarray, document: str):
62
+ item_id = len(self.documents)
63
+ self.index.add_item(item_id, embedding)
64
+ self.documents.append(document)
65
+
66
+ def build(self, num_trees:int = 10):
67
+ self.index.build(num_trees)
68
+
69
+ def save(self, filepath: str):
70
+ self.index.save(filepath)
71
+
72
+ def load(self, filepath: str):
73
+ self.index.load(filepath)
74
+
75
+ def search(self, query_embedding: np.ndarray, top_k: int = 5) -> List[Dict[str, Union[str, float]]]:
76
+ indices, distances = self.index.get_nns_by_vector(
77
+ query_embedding, top_k, include_distances=True
78
+ )
79
+
80
+ results = [
81
+ {
82
+ "document": self.documents[idx],
83
+ "score": 1 / (1 + distance) # Convert distance to similarity
84
+ } for idx, distance in zip(indices, distances)
85
+ ]
86
+
87
+ return results
88
+
89
+
SmartSearch/database/chromadb.py ADDED
@@ -0,0 +1,213 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, Any, Optional, List
2
+
3
+ import chromadb
4
+ from chromadb.config import Settings
5
+ from chromadb.api.types import (
6
+ Where,
7
+ GetResult,
8
+ QueryResult,
9
+ )
10
+
11
+ from ..embedding_provider import EmbeddingProvider
12
+ from .vector_store import VectorStore
13
+
14
+
15
+ class ChromaDB(VectorStore):
16
+ """
17
+ ChromaDB is an example of a vector-store class implementation.
18
+
19
+ See more:
20
+ https://github.com/chroma-core/chroma
21
+ """
22
+
23
+ def __init__(
24
+ self,
25
+ configs: Dict[str, Any] = {},
26
+ db_path: str = ".chromadb",
27
+ embedding_function: Optional[EmbeddingProvider] = None,
28
+ collection_name: Optional[str] = None,
29
+ ) -> None:
30
+ self.client = chromadb.PersistentClient(
31
+ path=db_path
32
+ )
33
+ self.configs = configs
34
+
35
+ self.embedding_function = embedding_function
36
+ self._collection_name = collection_name
37
+
38
+ self.collection = self.client.get_or_create_collection(
39
+ name = self.collection_name or "default_collection"
40
+ )
41
+
42
+ # self.logger = get_logger(self.__class__.__name__)
43
+
44
+ @property
45
+ def db_path(self) -> str:
46
+ return self.client.get_settings().persist_directory
47
+
48
+ @db_path.setter
49
+ def db_path(self, value: str) -> None:
50
+ self.client = chromadb.PersistentClient(path=value)
51
+
52
+ self.collection = self.client.get_or_create_collection(
53
+ name = self.collection_name or "default_collection"
54
+ )
55
+
56
+ @property
57
+ def collection_name(self):
58
+ return self._collection_name
59
+
60
+ @collection_name.setter
61
+ def collection_name(self, value):
62
+ self._collection_name = value
63
+ self.collection.modify(name=value)
64
+
65
+ def add_data(
66
+ self,
67
+ documents: List[str],
68
+ ids: List[str],
69
+ metadatas: Optional[List[Dict[str, Any]]] = None,
70
+ **optional_kwargs
71
+ ) -> None:
72
+ """
73
+ Add data to the collection by creating embeddings for them.
74
+
75
+ Args:
76
+ documents (List[str]): List of documents to add.
77
+ ids (List[str]): List of ids for the documents.
78
+ metadatas (Optional[List[Dict[str, Any]]]): List of metadata for the documents.
79
+ **optional_kwargs: Additional keyword arguments (see collection.add for more).
80
+ """
81
+
82
+ try:
83
+ params = {
84
+ "documents": documents,
85
+ "ids": ids,
86
+ **optional_kwargs
87
+ }
88
+
89
+ params["metadatas"] = metadatas or None
90
+
91
+ # If an embedding function is provided, create embeddings for the documents
92
+ if self.embedding_function:
93
+ embeddings = self.embedding_function.embed_documents(documents)
94
+ params["embeddings"] = embeddings
95
+
96
+ self.collection.add(**params)
97
+ except Exception as e:
98
+ # self.logger.error(f"Error adding data to collection: {e}")
99
+ print(f"Error adding data to collection: {e}")
100
+ raise e
101
+
102
+ def search(
103
+ self,
104
+ query_text: Optional[List[str]] = None,
105
+ query_embedding: Optional[List[List[float]]] = None,
106
+ n_results: int = 10,
107
+ **optional_kwargs
108
+ ) -> QueryResult:
109
+ """
110
+ Query the collection for similar documents.
111
+
112
+ Args:
113
+ query_text (Optional[List[str]]): List of query texts.
114
+ query_embedding (Optional[List[List[float]]]): List of query embeddings.
115
+ n_results (int): Number of results to return.
116
+ **optional_kwargs: Additional keyword arguments (see collection.query for more).
117
+
118
+ Returns:
119
+ QueryResult: The result of the query.
120
+ """
121
+
122
+ try:
123
+ if query_text is None and query_embedding is None:
124
+ raise ValueError("Either query_text or query_embedding must be provided.")
125
+
126
+ params = {
127
+ "n_results": n_results,
128
+ **optional_kwargs
129
+ }
130
+
131
+ if query_text and query_embedding is None:
132
+ if self.embedding_function:
133
+ query_embedding = self.embedding_function.embed_query(query_text)
134
+ params["query_embeddings"] = query_embedding
135
+ else:
136
+ params["query_text"] = query_text
137
+
138
+ elif query_embedding and query_text is None:
139
+ params["query_embeddings"] = query_embedding
140
+
141
+ elif query_embedding and query_text:
142
+ params["query_embeddings"] = query_embedding
143
+
144
+ if self.embedding_function:
145
+ embeddings = self.embedding_function.embed_query(query_text)
146
+ params["query_embeddings"] = query_embedding.extend(embeddings)
147
+ else:
148
+ params["query_text"] = query_text
149
+
150
+ return self.collection.query(**params)
151
+ except Exception as e:
152
+ # self.logger.error(f"Error querying data from collection: {e}")
153
+ print(f"Error querying data from collection: {e}")
154
+ raise e
155
+
156
+ def query_by_id_or_metadata(
157
+ self,
158
+ ids: Optional[List[str]] = None,
159
+ where: Optional[Where] = None,
160
+ n_results: int = 10,
161
+ **optional_kwargs
162
+ ) -> GetResult:
163
+ """
164
+ Query the collection for similar documents.
165
+
166
+ Args:
167
+ ids (Optional[List[str]]): List of ids to query.
168
+ where (Optional[Where]): Where clause to query.
169
+ n_results (int): Number of results to return.
170
+ **optional_kwargs: Additional keyword arguments (see collection.get for more).
171
+
172
+ Returns:
173
+ GetResult: The result of the query.
174
+ """
175
+
176
+ try:
177
+ if ids is None and where is None:
178
+ raise ValueError("Either ids or where must be provided.")
179
+
180
+ params = {
181
+ "n_results": n_results,
182
+ **optional_kwargs
183
+ }
184
+
185
+ if ids:
186
+ params["ids"] = ids
187
+ if where:
188
+ params["where"] = where
189
+
190
+ return self.collection.get(**params)
191
+ except Exception as e:
192
+ # self.logger.error(f"Error querying data from collection: {e}")
193
+ print(f"Error querying data from collection: {e}")
194
+ raise e
195
+
196
+ def delete_collection(self, collection_name: Optional[str] = None) -> None:
197
+ """
198
+ Delete a specific collection from the ChromaDB.
199
+
200
+ Args:
201
+ collection_name (Optional[str]): Name of collection to delete.
202
+ Uses class's collection_name if not provided.
203
+ """
204
+ try:
205
+ target_collection = collection_name or self.collection_name
206
+ if not target_collection:
207
+ raise ValueError("No collection name provided")
208
+
209
+ self.client.delete_collection(name=target_collection)
210
+ print(f"Collection '{target_collection}' deleted successfully.")
211
+ except Exception as e:
212
+ print(f"Error deleting collection: {e}")
213
+
SmartSearch/database/vector_store.py ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from abc import ABC, abstractmethod
2
+
3
+ from typing import Any
4
+
5
+ class VectorStore(ABC):
6
+ @abstractmethod
7
+ def add_data(self, *args, **kwargs) -> Any:
8
+ """
9
+ Add data to the vector store
10
+ """
11
+ pass
12
+
13
+ @abstractmethod
14
+ def search(self, *args, **kwargs) -> Any:
15
+ """
16
+ Search data from vector store
17
+ """
18
+ pass
SmartSearch/embedding_provider.py ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from abc import ABC, abstractmethod
2
+ from typing import List
3
+ import numpy as np
4
+
5
+ class EmbeddingProvider(ABC):
6
+ """
7
+ Abstract class for the llm providers
8
+ """
9
+ @abstractmethod
10
+ def embed_documents(self, documents: List[str]) -> np.ndarray:
11
+ """Embed a list of documents"""
12
+ pass
13
+
14
+ @abstractmethod
15
+ def embed_query(self, query: str) -> np.ndarray:
16
+ """Embed a query"""
17
+ pass
18
+
SmartSearch/hybrid_search.py ADDED
@@ -0,0 +1,102 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from typing import List, Dict, Union, Optional, Any
3
+ import numpy as np
4
+
5
+ from .embedding_provider import EmbeddingProvider
6
+ from .database.annoydb import AnnoyDB
7
+ from .keyword_search_provider import KeywordSearchProvider
8
+
9
+ class HybridSearch:
10
+ def __init__(
11
+ self,
12
+ embedding_provider: EmbeddingProvider,
13
+ documents: List[str] = None,
14
+ ann_filepath: Optional[str] = None,
15
+ semantic_weight: float = 0.7,
16
+ keyword_weight: float = 0.3
17
+ ) -> None:
18
+ self.embedding_provider = embedding_provider
19
+ self.documents = documents
20
+
21
+ if ann_filepath and os.path.exists(ann_filepath):
22
+ self.index = AnnoyDB
23
+ self.embeddings = self.embedding_provider.embed_documents(documents)
24
+
25
+ self.vector_db = AnnoyDB(
26
+ embedding_dim=self.embeddings.shape[1]
27
+ )
28
+
29
+ for emb, doc in zip(self.embeddings, documents):
30
+ self.vector_db.add_data(emb, doc)
31
+ self.vector_db.build()
32
+
33
+ # Keyword Search Setup
34
+ self.keyword_search = KeywordSearchProvider(documents)
35
+
36
+ # Weights for hybrid search
37
+ self.semantic_weight = semantic_weight
38
+ self.keyword_weight = keyword_weight
39
+
40
+ self.documents = documents
41
+
42
+ def hybrid_search(self, query: str, top_k: int = 5) -> List[Dict[str, Union[str, float]]]:
43
+ # Embed query
44
+ query_embedding = self.embedding_provider.embed_query(query)
45
+
46
+ # Perform semantic search
47
+ semantic_results = self.vector_db.search(query_embedding, top_k)
48
+
49
+ # Perform keyword search
50
+ keyword_results = self.keyword_search.search(query, top_k)
51
+
52
+ # Combine results with weighted scoring
53
+ combined_results = {}
54
+
55
+ for result in semantic_results:
56
+ doc = result['document']
57
+ combined_results[doc] = {
58
+ 'semantic_score': result['score'] * self.semantic_weight,
59
+ 'keyword_score': 0,
60
+ 'hybrid_score': result['score'] * self.semantic_weight
61
+ }
62
+
63
+ for result in keyword_results:
64
+ doc = result['document']
65
+ if doc in combined_results:
66
+ combined_results[doc]['keyword_score'] = result['score'] * self.keyword_weight
67
+ combined_results[doc]['hybrid_score'] += result['score'] * self.keyword_weight
68
+ else:
69
+ combined_results[doc] = {
70
+ 'semantic_score': 0,
71
+ 'keyword_score': result['score'] * self.keyword_weight,
72
+ 'hybrid_score': result['score'] * self.keyword_weight
73
+ }
74
+
75
+ # Sort and return top results
76
+ sorted_results = sorted(
77
+ [
78
+ {**{'document': doc}, **scores}
79
+ for doc, scores in combined_results.items()
80
+ ],
81
+ key=lambda x: x['hybrid_score'],
82
+ reverse=True
83
+ )
84
+
85
+ return sorted_results[:top_k]
86
+
87
+ def set_weights(self, semantic_weight: float, keyword_weight: float):
88
+ """
89
+ Dynamically update search weights.
90
+
91
+ Args:
92
+ semantic_weight: New weight for semantic search
93
+ keyword_weight: New weight for keyword search
94
+ """
95
+ if not (0 <= semantic_weight <= 1 and 0 <= keyword_weight <= 1):
96
+ raise ValueError("Weights must be between 0 and 1")
97
+
98
+ if not np.isclose(semantic_weight + keyword_weight, 1.0):
99
+ raise ValueError("Semantic and keyword weights must sum to 1.0")
100
+
101
+ self.semantic_weight = semantic_weight
102
+ self.keyword_weight = keyword_weight
SmartSearch/keyword_search_provider.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Dict, Union
2
+
3
+ class KeywordSearchProvider:
4
+ def __init__(self, documents: List[str]):
5
+ from sklearn.feature_extraction.text import TfidfVectorizer
6
+ self.vectorizer = TfidfVectorizer()
7
+ self.tfidf_matrix = self.vectorizer.fit_transform(documents)
8
+ self.documents = documents
9
+
10
+ def search(self, query: str, top_k: int = 5) -> List[Dict[str, Union[str, float]]]:
11
+ from sklearn.metrics.pairwise import cosine_similarity
12
+ query_vector = self.vectorizer.transform([query])
13
+ similarities = cosine_similarity(query_vector, self.tfidf_matrix)[0]
14
+
15
+ # Get top-k results
16
+ top_indices = similarities.argsort()[-top_k:][::-1]
17
+ results = [
18
+ {
19
+ "document": self.documents[idx],
20
+ "score": similarities[idx]
21
+ } for idx in top_indices
22
+ ]
23
+
24
+ return results
SmartSearch/providers/OpenAIEmbedding.py ADDED
@@ -0,0 +1,125 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Dict, Union, Optional
2
+ import tiktoken
3
+ from ..embedding_provider import EmbeddingProvider
4
+ import numpy as np
5
+
6
+ class OpenAIEmbedding(EmbeddingProvider):
7
+ def __init__(
8
+ self,
9
+ api_key: Optional[str] = None,
10
+ model: str = "text-embedding-3-small",
11
+ max_tokens: int = 8191
12
+ ) -> None:
13
+ """Initialize OpenAI embedding provider
14
+
15
+ Args:
16
+ model_name (str, optional): Name of the embedding model. Default to "text-embedding-3-small"
17
+ more info: https://platform.openai.com/docs/models#embeddings
18
+ api_key: api_key for OpenAI
19
+ """
20
+ from openai import OpenAI
21
+
22
+ self.client = OpenAI(api_key=api_key)
23
+ self.model = model
24
+ self.max_tokens = max_tokens
25
+ self.tokenizer = tiktoken.encoding_for_model(model)
26
+
27
+ def _trancated_text(self, text: str) -> str:
28
+ """Truncate text into maximum token length
29
+
30
+ Args:
31
+ text (str): Input text
32
+
33
+ Returns:
34
+ str: Truncated text
35
+ """
36
+ tokens = self.tokenizer.encode(text)
37
+ truncated_tokens = tokens[:self.max_tokens]
38
+ return self.tokenizer.decode(truncated_tokens)
39
+
40
+ def embed_documents(
41
+ self,
42
+ documents: List[str],
43
+ batch_size: int = 100
44
+ ) -> np.array:
45
+ """Embed a list of documents
46
+
47
+ Args:
48
+ documents (List[str]): List of documents to embed
49
+
50
+ Returns:
51
+ np.array: embeddings of documents
52
+ """
53
+ truncated_docs = [self._trancated_text(doc) for doc in documents]
54
+
55
+ embeddings = []
56
+ for i in range(0, len(truncated_docs), batch_size):
57
+ batch = truncated_docs[i: i+batch_size]
58
+
59
+ response = self.client.embeddings.create(
60
+ input=batch,
61
+ model=self.model
62
+ )
63
+ batch_embeddings = [
64
+ embed.embedding for embed in response.data
65
+ ]
66
+ embeddings.extend(batch_embeddings)
67
+
68
+ return np.array(embeddings)
69
+
70
+ def embed_query(self, query):
71
+ truncated_query = self._trancated_text(query)
72
+
73
+ response = self.client.embeddings.create(
74
+ input=[truncated_query],
75
+ model=self.model
76
+ )
77
+ return np.array(response.data[0].embedding)
78
+
79
+ def get_embedding_info(self) -> Dict[str, Union[str, int]]:
80
+ """
81
+ Get information about the current embedding configuration
82
+
83
+ Returns:
84
+ Dict: Embedding configuration details
85
+ """
86
+ return {
87
+ "model": self.model,
88
+ "max_tokens": self.max_tokens,
89
+ "batch_size": 100, # Default batch size
90
+ }
91
+
92
+ def list_available_models(self) -> List[str]:
93
+ """
94
+ List available OpenAI embedding models
95
+
96
+ Returns:
97
+ List[str]: Available embedding model names
98
+ """
99
+ return [
100
+ "text-embedding-ada-002", # Most common
101
+ "text-embedding-3-small", # Newer, more efficient
102
+ "text-embedding-3-large" # Highest quality
103
+ ]
104
+
105
+ def estimate_cost(self, num_documents: int) -> float:
106
+ """
107
+ Estimate embedding cost
108
+
109
+ Args:
110
+ num_documents (int): Number of documents to embed
111
+
112
+ Returns:
113
+ float: Estimated cost in USD
114
+ """
115
+ # Pricing as of 2024 (subject to change)
116
+ pricing = {
117
+ "text-embedding-ada-002": 0.0001 / 1000, # $0.0001 per 1000 tokens
118
+ "text-embedding-3-small": 0.00006 / 1000,
119
+ "text-embedding-3-large": 0.00013 / 1000
120
+ }
121
+
122
+ # Estimate tokens (assuming ~100 tokens per document)
123
+ total_tokens = num_documents * 100
124
+
125
+ return total_tokens * pricing.get(self.model, pricing["text-embedding-ada-002"])
SmartSearch/providers/SentenceTransformerEmbedding.py ADDED
@@ -0,0 +1,85 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Dict, Union
2
+ from ..embedding_provider import EmbeddingProvider
3
+ import numpy as np
4
+
5
+ class SentenceTransformerEmbedding(EmbeddingProvider):
6
+ def __init__(
7
+ self,
8
+ model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
9
+ device: str = None,
10
+ batch_size: int = 32,
11
+ normalize_embeddings: bool = True
12
+ ) -> None:
13
+ """Initialize sentence transformer embedding provider
14
+
15
+ Args:
16
+ model_name (str, optional): Name of the sentence tranformer model. Defaults to "sentence-transformers/all-MiniLM-L6-v2".
17
+ """
18
+ from sentence_transformers import SentenceTransformer
19
+
20
+ if device is None:
21
+ import torch
22
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
23
+
24
+ self.model = SentenceTransformer(model_name, device=device)
25
+ self.model_name = model_name
26
+ self.batch_size = batch_size
27
+ self.normalize_embeddings = normalize_embeddings
28
+
29
+ def embed_documents(self, documents: List[str]) -> np.ndarray:
30
+ """Embed a list of documents
31
+
32
+ Args:
33
+ documents (List[str]): List of documents to embed
34
+ """
35
+ return self.model.encode(
36
+ documents,
37
+ batch_size=self.batch_size,
38
+ normalize_embeddings=self.normalize_embeddings
39
+ )
40
+
41
+ def embed_query(self, query: str) -> np.ndarray:
42
+ """Embed a single query
43
+
44
+ Args:
45
+ query (str): Query to embed
46
+
47
+ Returns:
48
+ np.ndarray: Embedding vector
49
+ """
50
+ return self.model.encode(
51
+ query,
52
+ normalize_embeddings=self.normalize_embeddings
53
+ )
54
+
55
+ def get_model_info(self) -> Dict[str, Union[str, int]]:
56
+ """
57
+ Retrieve information about the current embedding model
58
+
59
+ Returns:
60
+ Dict: Model information
61
+ """
62
+ return {
63
+ "model_name": self.model_name,
64
+ "device": self.device,
65
+ "batch_size": self.batch_size,
66
+ "normalize_embeddings": self.normalize_embeddings,
67
+ "embedding_dim": self.model.get_sentence_embedding_dimension()
68
+ }
69
+
70
+ def list_available_models(self) -> List[str]:
71
+ """
72
+ List some popular Sentence Transformer models
73
+
74
+ Returns:
75
+ List[str]: Available model names
76
+ """
77
+ popular_models = [
78
+ "sentence-transformers/all-MiniLM-L6-v2", # Small and fast
79
+ "sentence-transformers/all-mpnet-base-v2", # High performance
80
+ "sentence-transformers/all-distilroberta-v1", # Lightweight
81
+ "sentence-transformers/multi-qa-MiniLM-L6-cos-v1", # Question Answering
82
+ "sentence-transformers/multi-qa-mpnet-base-cos-v1", # Multilingual QA
83
+ "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2" # Multilingual
84
+ ]
85
+ return popular_models
SmartSearch/providers/__init__.py ADDED
File without changes
SmartSearch/search_manager.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List
2
+
3
+ from embedding_provider import EmbeddingProvider
4
+ from database.annoydb import AnnoyDB
5
+
6
+ class SearchManager:
7
+ def __init__(
8
+ self,
9
+ embedding_provider: EmbeddingProvider,
10
+ documents: List[str],
11
+ semantic_weight: float = 0.7,
12
+ keyword_weight: float = 0.3
13
+ ) -> None:
14
+ """Smart Search Manager
15
+
16
+ Args:
17
+ embedding_provider (EmbeddingProvider): embedding provider
18
+ documents (List[str]): list of documents
19
+ semantic_weight (float, optional): _description_. Defaults to 0.7.
20
+ keyword_weight (float, optional): _description_. Defaults to 0.3.
21
+ """
22
+ self.embedding_provider = embedding_provider
23
+ self.semantic_embeddings = embedding_provider.embed_documents(documents)
24
+
25
+ # Vector Database Setup
26
+ self.vector_db = AnnoyDB(
27
+ embedding_dim=self.semantic_embeddings.shape[1]
28
+ )
29
+ for emb, doc in zip(self.semantic_embeddings, documents):
30
+ self.vector_db.add_item(emb, doc)
31
+ self.vector_db.build()
32
+
33
+ # Keyword Search Setup
34
+ self.keyword_search = KeywordSearchProvider(documents)
35
+
36
+ # Weights for hybrid search
37
+ self.semantic_weight = semantic_weight
38
+ self.keyword_weight = keyword_weight
39
+
40
+ self.documents = documents
app.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List
2
+ from dataclasses import asdict
3
+ import pandas as pd
4
+ import gradio as gr
5
+
6
+ from SmartSearch.database.chromadb import ChromaDB
7
+ from SmartSearch.providers.SentenceTransformerEmbedding import SentenceTransformerEmbedding
8
+ from utils import combine_metadata_with_distance
9
+ st_chroma = ChromaDB(
10
+ embedding_function=SentenceTransformerEmbedding(model_name='all-mpnet-base-v2'),
11
+ collection_name="novel_mockup_collection"
12
+ )
13
+
14
+ # Function to search for products
15
+ def search_novels(query, k):
16
+ result = st_chroma.search(query_text=query, n_results=k)
17
+
18
+ result = combine_metadata_with_distance(result['metadatas'], result['distances'])
19
+ result = pd.DataFrame(result)
20
+ return result
21
+
22
+ with gr.Blocks() as demo:
23
+ with gr.Row():
24
+ query = gr.Textbox(label="Search Query", placeholder="write a query to find the courses")
25
+ with gr.Row():
26
+ # search_type = gr.Dropdown(label="Search Type", choices=['semantic', 'keyword', 'hybrid'], value='hybrid')
27
+ k = gr.Number(label="Items Count", value=10)
28
+ # rerank = gr.Checkbox(value=True, label="Rerank")
29
+ results = gr.Dataframe(label="Search Results")
30
+
31
+ search_button = gr.Button("Search", variant='primary')
32
+ search_button.click(fn=search_novels, inputs=[query, k], outputs=results)
33
+
34
+ demo.launch()
requirements.txt ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ -i https://pypi.org/simple
2
+ aiofiles==23.2.1; python_version >= '3.7'
3
+ annotated-types==0.7.0; python_version >= '3.8'
4
+ anyio==4.8.0; python_version >= '3.9'
5
+ asgiref==3.8.1; python_version >= '3.8'
6
+ backoff==2.2.1; python_version >= '3.7' and python_version < '4.0'
7
+ bcrypt==4.2.1; python_version >= '3.7'
8
+ build==1.2.2.post1; python_version >= '3.8'
9
+ cachetools==5.5.0; python_version >= '3.7'
10
+ certifi==2024.12.14; python_version >= '3.6'
11
+ charset-normalizer==3.4.1; python_version >= '3.7'
12
+ chroma-hnswlib==0.7.6
13
+ chromadb==0.6.2; python_version >= '3.9'
14
+ click==8.1.8; python_version >= '3.7'
15
+ coloredlogs==15.0.1; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3, 3.4'
16
+ deprecated==1.2.15; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3'
17
+ durationpy==0.9
18
+ exceptiongroup==1.2.2; python_version >= '3.7'
19
+ fastapi==0.115.6; python_version >= '3.8'
20
+ ffmpy==0.5.0; python_version >= '3.8' and python_version < '4.0'
21
+ filelock==3.16.1; python_version >= '3.8'
22
+ flatbuffers==24.12.23
23
+ fsspec==2024.12.0; python_version >= '3.8'
24
+ google-auth==2.37.0; python_version >= '3.7'
25
+ googleapis-common-protos==1.66.0; python_version >= '3.7'
26
+ gradio==5.11.0; python_version >= '3.10'
27
+ gradio-client==1.5.3; python_version >= '3.10'
28
+ grpcio==1.69.0; python_version >= '3.8'
29
+ h11==0.14.0; python_version >= '3.7'
30
+ httpcore==1.0.7; python_version >= '3.8'
31
+ httptools==0.6.4; python_full_version >= '3.8.0'
32
+ httpx==0.28.1; python_version >= '3.8'
33
+ huggingface-hub==0.27.1; python_full_version >= '3.8.0'
34
+ humanfriendly==10.0; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2, 3.3, 3.4'
35
+ idna==3.10; python_version >= '3.6'
36
+ importlib-metadata==8.5.0; python_version >= '3.8'
37
+ importlib-resources==6.5.2; python_version >= '3.9'
38
+ jinja2==3.1.5; python_version >= '3.7'
39
+ joblib==1.4.2; python_version >= '3.8'
40
+ kubernetes==31.0.0; python_version >= '3.6'
41
+ markdown-it-py==3.0.0; python_version >= '3.8'
42
+ markupsafe==2.1.5; python_version >= '3.7'
43
+ mdurl==0.1.2; python_version >= '3.7'
44
+ mmh3==5.0.1; python_version >= '3.8'
45
+ monotonic==1.6
46
+ mpmath==1.3.0
47
+ networkx==3.4.2; python_version >= '3.10'
48
+ numpy==2.2.1; python_version >= '3.10'
49
+ oauthlib==3.2.2; python_version >= '3.6'
50
+ onnxruntime==1.20.1
51
+ opentelemetry-api==1.29.0; python_version >= '3.8'
52
+ opentelemetry-exporter-otlp-proto-common==1.29.0; python_version >= '3.8'
53
+ opentelemetry-exporter-otlp-proto-grpc==1.29.0; python_version >= '3.8'
54
+ opentelemetry-instrumentation==0.50b0; python_version >= '3.8'
55
+ opentelemetry-instrumentation-asgi==0.50b0; python_version >= '3.8'
56
+ opentelemetry-instrumentation-fastapi==0.50b0; python_version >= '3.8'
57
+ opentelemetry-proto==1.29.0; python_version >= '3.8'
58
+ opentelemetry-sdk==1.29.0; python_version >= '3.8'
59
+ opentelemetry-semantic-conventions==0.50b0; python_version >= '3.8'
60
+ opentelemetry-util-http==0.50b0; python_version >= '3.8'
61
+ orjson==3.10.14; python_version >= '3.8'
62
+ overrides==7.7.0; python_version >= '3.6'
63
+ packaging==24.2; python_version >= '3.8'
64
+ pandas==2.2.3; python_version >= '3.9'
65
+ pillow==11.1.0; python_version >= '3.9'
66
+ posthog==3.7.5
67
+ protobuf==5.29.3; python_version >= '3.8'
68
+ pyasn1==0.6.1; python_version >= '3.8'
69
+ pyasn1-modules==0.4.1; python_version >= '3.8'
70
+ pydantic==2.10.4; python_version >= '3.8'
71
+ pydantic-core==2.27.2; python_version >= '3.8'
72
+ pydub==0.25.1
73
+ pygments==2.19.1; python_version >= '3.8'
74
+ pypika==0.48.9
75
+ pyproject-hooks==1.2.0; python_version >= '3.7'
76
+ python-dateutil==2.9.0.post0; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2'
77
+ python-dotenv==1.0.1; python_version >= '3.8'
78
+ python-multipart==0.0.20; python_version >= '3.8'
79
+ pytz==2024.2
80
+ pyyaml==6.0.2; python_version >= '3.8'
81
+ regex==2024.11.6; python_version >= '3.8'
82
+ requests==2.32.3; python_version >= '3.8'
83
+ requests-oauthlib==2.0.0; python_version >= '3.4'
84
+ rich==13.9.4; python_full_version >= '3.8.0'
85
+ rsa==4.9; python_version >= '3.6' and python_version < '4'
86
+ ruff==0.8.6; python_version >= '3.7'
87
+ safehttpx==0.1.6; python_version >= '3.10'
88
+ safetensors==0.5.2; python_version >= '3.7'
89
+ scikit-learn==1.6.0; python_version >= '3.9'
90
+ scipy==1.15.0; python_version >= '3.10'
91
+ semantic-version==2.10.0; python_version >= '2.7'
92
+ sentence-transformers==3.3.1; python_version >= '3.9'
93
+ shellingham==1.5.4; python_version >= '3.7'
94
+ six==1.17.0; python_version >= '2.7' and python_version not in '3.0, 3.1, 3.2'
95
+ sniffio==1.3.1; python_version >= '3.7'
96
+ starlette==0.41.3; python_version >= '3.8'
97
+ sympy==1.13.1; python_version >= '3.8'
98
+ tenacity==9.0.0; python_version >= '3.8'
99
+ threadpoolctl==3.5.0; python_version >= '3.8'
100
+ tokenizers==0.21.0; python_version >= '3.7'
101
+ tomli==2.2.1; python_version >= '3.8'
102
+ tomlkit==0.13.2; python_version >= '3.8'
103
+ torch==2.5.1; python_full_version >= '3.8.0'
104
+ tqdm==4.67.1; python_version >= '3.7'
105
+ transformers==4.47.1; python_full_version >= '3.9.0'
106
+ typer==0.15.1; python_version >= '3.7'
107
+ typing-extensions==4.12.2; python_version >= '3.8'
108
+ tzdata==2024.2; python_version >= '2'
109
+ urllib3==2.3.0; python_version >= '3.9'
110
+ uvicorn[standard]==0.34.0; python_version >= '3.9'
111
+ uvloop==0.21.0; python_full_version >= '3.8.0'
112
+ watchfiles==1.0.3; python_version >= '3.9'
113
+ websocket-client==1.8.0; python_version >= '3.8'
114
+ websockets==14.1; python_version >= '3.9'
115
+ wrapt==1.17.0; python_version >= '3.8'
116
+ zipp==3.21.0; python_version >= '3.9'
utils.py ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ def combine_metadata_with_distance(metadatas, distances):
2
+ # Flatten the nested lists if they are nested
3
+ metadatas = metadatas[0] if isinstance(metadatas[0], list) else metadatas
4
+ distances = distances[0] if isinstance(distances[0], list) else distances
5
+ print(metadatas)
6
+ if len(metadatas) != len(distances):
7
+ raise ValueError("Number of metadata entries must match the number of distances")
8
+
9
+ combined_result = []
10
+ for metadata, distance in zip(metadatas, distances):
11
+ new_metadata = {
12
+ 'title': metadata.get('title', ''),
13
+ 'description': metadata.get('description', ''),
14
+ 'price': metadata.get('price', ''),
15
+ 'totalRatings': metadata.get('totalRatings', 0),
16
+ 'reviewSummary': metadata.get('reviewSummary', ''),
17
+ 'triggerWarning': metadata.get('triggerWarning', ''),
18
+ 'distance': distance
19
+ }
20
+ combined_result.append(new_metadata)
21
+
22
+ combined_result.sort(key=lambda x: x['distance'])
23
+
24
+ return combined_result