simonbb 0.0.9__tar.gz → 0.0.11__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {simonbb-0.0.9 → simonbb-0.0.11}/PKG-INFO +22 -12
- simonbb-0.0.11/pyproject.toml +44 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/bert.py +3 -3
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis438/rag.py +205 -113
- simonbb-0.0.9/pyproject.toml +0 -34
- {simonbb-0.0.9 → simonbb-0.0.11}/LICENSE +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/README.md +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/__init__.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/__init__.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/actions.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/agent.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/dqn.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/experience.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/gpt.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/ignite.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/math500.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/qwen3.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/transformer.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/utils.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis433/wrappers.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/cis438/__init__.py +0 -0
- {simonbb-0.0.9 → simonbb-0.0.11}/src/simonbb/utils.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: simonbb
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.11
|
|
4
4
|
Summary: A small package to facilitate teaching AI courses at business schools
|
|
5
5
|
Project-URL: Homepage, https://simonbb.com
|
|
6
6
|
Author-email: Huaxia Rui <huaxia.rui@simon.rochester.edu>
|
|
@@ -9,17 +9,27 @@ License-File: LICENSE
|
|
|
9
9
|
Classifier: Operating System :: OS Independent
|
|
10
10
|
Classifier: Programming Language :: Python :: 3
|
|
11
11
|
Requires-Python: >=3.9
|
|
12
|
-
Requires-Dist:
|
|
13
|
-
Requires-Dist:
|
|
14
|
-
Requires-Dist:
|
|
15
|
-
Requires-Dist:
|
|
16
|
-
Requires-Dist:
|
|
17
|
-
Requires-Dist:
|
|
18
|
-
Requires-Dist:
|
|
19
|
-
Requires-Dist:
|
|
20
|
-
Requires-Dist:
|
|
21
|
-
Requires-Dist:
|
|
22
|
-
Requires-Dist:
|
|
12
|
+
Requires-Dist: chromadb~=1.5.9
|
|
13
|
+
Requires-Dist: gymnasium~=1.3.0
|
|
14
|
+
Requires-Dist: langchain-chroma~=1.1.0
|
|
15
|
+
Requires-Dist: langchain-community~=0.4.2
|
|
16
|
+
Requires-Dist: langchain-core~=1.6.12
|
|
17
|
+
Requires-Dist: langchain-openai~=1.7.2
|
|
18
|
+
Requires-Dist: langchain-text-splitters~=1.1.3
|
|
19
|
+
Requires-Dist: langchain[mcp]~=1.4.6
|
|
20
|
+
Requires-Dist: langgraph~=1.2.16
|
|
21
|
+
Requires-Dist: matplotlib~=3.10.9
|
|
22
|
+
Requires-Dist: numpy~=2.4.6
|
|
23
|
+
Requires-Dist: opencv-python~=5.0.0.93
|
|
24
|
+
Requires-Dist: pydantic~=2.14.1
|
|
25
|
+
Requires-Dist: python-dotenv~=1.2.4
|
|
26
|
+
Requires-Dist: pytorch-ignite~=0.5.5
|
|
27
|
+
Requires-Dist: pyyaml~=6.0.3
|
|
28
|
+
Requires-Dist: ray~=2.56.1
|
|
29
|
+
Requires-Dist: stable-baselines3~=2.9.0
|
|
30
|
+
Requires-Dist: torchvision~=0.24.1
|
|
31
|
+
Requires-Dist: torch~=2.9.1
|
|
32
|
+
Requires-Dist: transformers~=4.57.6
|
|
23
33
|
Description-Content-Type: text/markdown
|
|
24
34
|
|
|
25
35
|
# Background
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling >= 1.26"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "simonbb"
|
|
7
|
+
version = "0.0.11"
|
|
8
|
+
authors = [
|
|
9
|
+
{ name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
|
|
10
|
+
]
|
|
11
|
+
description = "A small package to facilitate teaching AI courses at business schools"
|
|
12
|
+
readme = "README.md"
|
|
13
|
+
requires-python = ">=3.9"
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Operating System :: OS Independent",
|
|
17
|
+
]
|
|
18
|
+
license = "MIT"
|
|
19
|
+
license-files = ["LICEN[CS]E*"]
|
|
20
|
+
dependencies = [
|
|
21
|
+
"matplotlib ~= 3.10.9",
|
|
22
|
+
"numpy ~= 2.4.6",
|
|
23
|
+
"torch ~= 2.9.1",
|
|
24
|
+
"torchvision ~= 0.24.1",
|
|
25
|
+
"transformers ~= 4.57.6",
|
|
26
|
+
"gymnasium ~= 1.3.0",
|
|
27
|
+
"pytorch-ignite ~= 0.5.5",
|
|
28
|
+
"stable_baselines3 ~= 2.9.0",
|
|
29
|
+
"ray ~= 2.56.1",
|
|
30
|
+
"opencv-python ~= 5.0.0.93",
|
|
31
|
+
"pyyaml ~= 6.0.3",
|
|
32
|
+
"langchain-core ~= 1.6.12",
|
|
33
|
+
"langchain[mcp] ~= 1.4.6",
|
|
34
|
+
"langgraph ~= 1.2.16",
|
|
35
|
+
"langchain-community ~= 0.4.2",
|
|
36
|
+
"langchain-chroma ~= 1.1.0",
|
|
37
|
+
"langchain-openai ~= 1.7.2",
|
|
38
|
+
"langchain_text_splitters ~= 1.1.3",
|
|
39
|
+
"chromadb ~= 1.5.9",
|
|
40
|
+
"python-dotenv ~= 1.2.4",
|
|
41
|
+
"pydantic ~= 2.14.1"
|
|
42
|
+
]
|
|
43
|
+
[project.urls]
|
|
44
|
+
Homepage = "https://simonbb.com"
|
|
@@ -1,11 +1,11 @@
|
|
|
1
1
|
import re
|
|
2
2
|
import torch
|
|
3
3
|
import torch.nn as nn
|
|
4
|
-
from transformers import
|
|
4
|
+
from transformers import AutoTokenizer, BertModel, BertConfig
|
|
5
5
|
|
|
6
6
|
#%%
|
|
7
7
|
|
|
8
|
-
tokenizer =
|
|
8
|
+
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
|
|
9
9
|
|
|
10
10
|
class BertClassifier( nn.Module ):
|
|
11
11
|
def __init__(self, freeze_bert=False):
|
|
@@ -33,7 +33,7 @@ def BertPreprocessor( txt ):
|
|
|
33
33
|
for s in txt:
|
|
34
34
|
s = re.sub(r'(@.*?)[\s]', ' ', s)
|
|
35
35
|
s = re.sub(r'\s+', ' ', s).strip()
|
|
36
|
-
encoded = tokenizer
|
|
36
|
+
encoded = tokenizer( text=s, add_special_tokens=True, max_length=128, padding='max_length', truncation=True, return_attention_mask=True, return_tensors="pt" )
|
|
37
37
|
input_ids.append( encoded.get('input_ids') ) # Add the outputs to the lists
|
|
38
38
|
attention_masks.append( encoded.get('attention_mask') )
|
|
39
39
|
return torch.tensor( input_ids ), torch.tensor( attention_masks )
|
|
@@ -1,8 +1,7 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
-
import os
|
|
4
3
|
import re
|
|
5
|
-
from typing import Annotated, TypedDict, get_args, Any,
|
|
4
|
+
from typing import Annotated, TypedDict, get_args, Any, Sequence
|
|
6
5
|
from dataclasses import dataclass
|
|
7
6
|
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
|
|
8
7
|
|
|
@@ -25,6 +24,7 @@ from langchain_community.document_loaders import PyPDFLoader, UnstructuredHTMLLo
|
|
|
25
24
|
from langchain_core.callbacks import CallbackManagerForRetrieverRun
|
|
26
25
|
|
|
27
26
|
import hashlib
|
|
27
|
+
from uuid import uuid4
|
|
28
28
|
import chromadb
|
|
29
29
|
from langchain_core.embeddings import Embeddings
|
|
30
30
|
from langchain_core.language_models import BaseChatModel
|
|
@@ -35,13 +35,10 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
35
35
|
|
|
36
36
|
def load_documents(path: Path):
|
|
37
37
|
"""Load documents from a file or directory."""
|
|
38
|
-
|
|
39
|
-
valid_extensions = {'.pdf', '.html', '.ejs'}
|
|
40
|
-
|
|
38
|
+
valid_extensions = {'.pdf', '.html', '.ejs'}
|
|
41
39
|
docs = []
|
|
42
40
|
success_cnt = 0
|
|
43
|
-
page_cnt = 0
|
|
44
|
-
|
|
41
|
+
page_cnt = 0
|
|
45
42
|
if path.exists():
|
|
46
43
|
if path.is_dir():
|
|
47
44
|
for file in [ item for item in path.rglob('*') if item.is_file() and item.suffix in valid_extensions ]:
|
|
@@ -77,25 +74,100 @@ def load_documents(path: Path):
|
|
|
77
74
|
|
|
78
75
|
return docs, success_cnt, page_cnt
|
|
79
76
|
|
|
77
|
+
#%% RAG Chains
|
|
78
|
+
|
|
79
|
+
system_prompt = """You are a helpful assistant answering questions based on provided document snippets.
|
|
80
|
+
IMPORTANT INSTRUCTIONS:
|
|
81
|
+
1. Base your answer ONLY on the provided context snippets below.
|
|
82
|
+
2. Synthesize information from multiple snippets if needed.
|
|
83
|
+
3. If the snippets contain partial information, provide the best answer you can with what's available.
|
|
84
|
+
4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
|
|
85
|
+
5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
|
|
86
|
+
|
|
87
|
+
The snippets are provided in this format:
|
|
88
|
+
Snippet 1: <text content>
|
|
89
|
+
Snippet 2: <text content>
|
|
90
|
+
...
|
|
91
|
+
|
|
92
|
+
{context}"""
|
|
93
|
+
|
|
94
|
+
def format_documents(documents: Sequence[Document]) -> str:
|
|
95
|
+
snippets = []
|
|
96
|
+
for i, doc in enumerate(documents, start=1):
|
|
97
|
+
snippets.append(f"Snippet {i}: {doc.page_content}")
|
|
98
|
+
context = "\n".join(snippets)
|
|
99
|
+
return context
|
|
100
|
+
|
|
101
|
+
def rag_basic(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
|
|
102
|
+
"""Return an LCEL chain accepting a question string and producing an answer string."""
|
|
103
|
+
prompt = ChatPromptTemplate.from_messages(
|
|
104
|
+
[
|
|
105
|
+
( "system", "Answer using only the context below. If the answer is absent, say you do not know.\n\nContext:\n{context}" ),
|
|
106
|
+
("human", "{query}"),
|
|
107
|
+
]
|
|
108
|
+
)
|
|
109
|
+
inputs = {
|
|
110
|
+
"context": retriever | format_documents,
|
|
111
|
+
"query": RunnablePassthrough(),
|
|
112
|
+
}
|
|
113
|
+
return inputs | prompt | model | StrOutputParser()
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
def rag_main(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
|
|
117
|
+
"""Return an LCEL chain accepting a question string and producing an dic with 4 keys/values
|
|
118
|
+
-- 'sources': original list of Documents from retriever
|
|
119
|
+
-- 'query': original query
|
|
120
|
+
-- 'context': merged content from relevant documents
|
|
121
|
+
-- 'answer': final answer """
|
|
122
|
+
prompt = ChatPromptTemplate.from_messages([
|
|
123
|
+
("system", system_prompt),
|
|
124
|
+
("human", "{query}")
|
|
125
|
+
])
|
|
126
|
+
return (
|
|
127
|
+
{
|
|
128
|
+
"sources": retriever,
|
|
129
|
+
"query": RunnablePassthrough(),
|
|
130
|
+
}
|
|
131
|
+
| RunnablePassthrough.assign( context=lambda x: format_documents(x["sources"]) )
|
|
132
|
+
| RunnablePassthrough.assign( answer=prompt | model | StrOutputParser() )
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
build_rag = rag_main
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
"""
|
|
139
|
+
Retrievers are subclasses of LangChain's BaseRetriever which inherits from Pydantic's BaseModel.
|
|
140
|
+
Unlike ordinary Python classes where attributes declared in the class body are class attributes shared by all instances,
|
|
141
|
+
Pydantic model fields declared with type annotations are instance attributes, not shared class attributes, even though they are declared in the class body.
|
|
142
|
+
"""
|
|
80
143
|
#%% MultiQueryRerankRetriever
|
|
81
144
|
|
|
82
145
|
class MultiQueryRerankRetriever(BaseRetriever):
|
|
83
146
|
"""Chroma retriever with optional multi-query retrieval and reranking."""
|
|
84
|
-
knowledge_base: str
|
|
85
147
|
embeddings: Embeddings = Field(exclude=True)
|
|
86
148
|
query_llm: BaseChatModel | None = Field(default=None, exclude=True)
|
|
87
149
|
reranker: Any = Field(default=None, exclude=True)
|
|
88
150
|
k: int = Field(default=3, ge=1)
|
|
89
151
|
top_n: int = Field(default=5, ge=1)
|
|
90
|
-
|
|
152
|
+
persist_directory: str | Path
|
|
153
|
+
collection_name: str
|
|
154
|
+
create_collection_if_not_exists: bool = False
|
|
155
|
+
store: Chroma = None
|
|
91
156
|
_base_retriever: BaseRetriever = PrivateAttr()
|
|
92
157
|
|
|
158
|
+
""" model_post_init() is a special method in Pydantic v2 that runs automatically after a model instance has been initialized and its fields have been validated."""
|
|
93
159
|
def model_post_init(self, __context: Any) -> None:
|
|
94
|
-
"""Initialize the underlying Chroma retriever."""
|
|
95
|
-
store = Chroma( persist_directory=
|
|
96
|
-
self._base_retriever = store.as_retriever( search_kwargs={"k": self.k} )
|
|
160
|
+
"""Initialize the underlying Chroma retriever. Use the last folder name as collection name"""
|
|
161
|
+
self.store = Chroma( persist_directory=self.persist_directory, embedding_function=self.embeddings, collection_name=self.collection_name, create_collection_if_not_exists=self.create_collection_if_not_exists )
|
|
162
|
+
self._base_retriever = self.store.as_retriever( search_kwargs={"k": self.k} )
|
|
97
163
|
if self.query_llm is not None and self.reranker is None:
|
|
98
164
|
raise ValueError( "A reranker is required when query_llm is provided." )
|
|
165
|
+
|
|
166
|
+
def ingest( self, documents: Sequence[Document], *, source_dir: str = 'placeholder', chunk_size: int = 1000, chunk_overlap: int = 150 ):
|
|
167
|
+
text_splitter = RecursiveCharacterTextSplitter( separators=["\n\n", "\n", ". ", " ", ""], chunk_size=chunk_size, chunk_overlap=chunk_overlap )
|
|
168
|
+
chunks = text_splitter.split_documents( documents )
|
|
169
|
+
self.store.add_documents( chunks )
|
|
170
|
+
|
|
99
171
|
# run_manager manages callbacks, tracing, and monitoring during a retriever's execution, allowing integrations with LangSmith
|
|
100
172
|
def _generate_queries( self, query: str, run_manager: CallbackManagerForRetrieverRun ) -> list[str]:
|
|
101
173
|
"""Generate five alternative formulations of the query."""
|
|
@@ -127,6 +199,7 @@ class MultiQueryRerankRetriever(BaseRetriever):
|
|
|
127
199
|
)
|
|
128
200
|
return [doc for _, doc in scored_docs[:self.top_n]]
|
|
129
201
|
|
|
202
|
+
# override _get_relevant_documents() method from BaseRetriever
|
|
130
203
|
def _get_relevant_documents( self, query: str, *, run_manager: CallbackManagerForRetrieverRun ) -> list[Document]:
|
|
131
204
|
"""Retrieve documents using simple or multi-query retrieval."""
|
|
132
205
|
# Simple retrieval
|
|
@@ -143,12 +216,91 @@ class MultiQueryRerankRetriever(BaseRetriever):
|
|
|
143
216
|
# Deduplicate and rerank
|
|
144
217
|
unique_docs = self._unique_union(doc_lists)
|
|
145
218
|
return self._rerank(query, unique_docs)
|
|
219
|
+
|
|
220
|
+
#%%
|
|
221
|
+
|
|
222
|
+
class DummyEmbeddings(Embeddings):
|
|
223
|
+
"""Return constant 1-dimensional embeddings without using an LLM."""
|
|
224
|
+
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
|
225
|
+
return [[0.0] for _ in texts]
|
|
226
|
+
def embed_query(self, text: str) -> list[float]:
|
|
227
|
+
return [0.0]
|
|
228
|
+
|
|
229
|
+
class ParentDocumentRetriever(BaseRetriever):
|
|
230
|
+
persist_directory: str | Path
|
|
231
|
+
parent_collection: str
|
|
232
|
+
child_collection: str
|
|
233
|
+
embeddings: Embeddings
|
|
234
|
+
k: int = 5
|
|
235
|
+
_parent_store: Chroma = PrivateAttr()
|
|
236
|
+
_child_store: Chroma = PrivateAttr()
|
|
237
|
+
|
|
238
|
+
def __init__(self, **kwargs):
|
|
239
|
+
super().__init__(**kwargs)
|
|
240
|
+
self._parent_store = Chroma(
|
|
241
|
+
embedding_function=DummyEmbeddings(),
|
|
242
|
+
collection_name=self.parent_collection,
|
|
243
|
+
persist_directory=self.persist_directory,
|
|
244
|
+
)
|
|
245
|
+
|
|
246
|
+
self._child_store = Chroma(
|
|
247
|
+
embedding_function=self.embeddings,
|
|
248
|
+
collection_name=self.child_collection,
|
|
249
|
+
persist_directory=self.persist_directory,
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
def ingest(self, docs: list[Document], *, source_dir: str = 'placeholder', parent_chunk_size: int = 600, parent_chunk_overlap: int = 100, child_chunk_size: int = 200, child_chunk_overlap: int = 50) -> None:
|
|
253
|
+
"""Split documents and store parent and child chunks."""
|
|
254
|
+
parent_splitter = RecursiveCharacterTextSplitter( chunk_size=parent_chunk_size, chunk_overlap=parent_chunk_overlap )
|
|
255
|
+
child_splitter = RecursiveCharacterTextSplitter( chunk_size=child_chunk_size, chunk_overlap=child_chunk_overlap )
|
|
256
|
+
|
|
257
|
+
parent_docs = []
|
|
258
|
+
parent_ids = []
|
|
259
|
+
child_docs = []
|
|
260
|
+
|
|
261
|
+
for parent in parent_splitter.split_documents(docs):
|
|
262
|
+
parent_id = str(uuid4())
|
|
263
|
+
parent_docs.append(parent)
|
|
264
|
+
parent_ids.append(parent_id)
|
|
265
|
+
for child in child_splitter.split_documents( [parent] ):
|
|
266
|
+
child.metadata["parent_id"] = parent_id
|
|
267
|
+
child_docs.append(child)
|
|
268
|
+
if parent_docs:
|
|
269
|
+
self._parent_store.add_documents( documents=parent_docs, ids=parent_ids )
|
|
270
|
+
if child_docs:
|
|
271
|
+
self._child_store.add_documents( documents=child_docs )
|
|
272
|
+
|
|
273
|
+
# override _get_relevant_documents() method from BaseRetriever
|
|
274
|
+
def _get_relevant_documents(
|
|
275
|
+
self,
|
|
276
|
+
query: str,
|
|
277
|
+
*,
|
|
278
|
+
run_manager: CallbackManagerForRetrieverRun,
|
|
279
|
+
) -> list[Document]:
|
|
280
|
+
"""Retrieve parent documents using child similarity search."""
|
|
281
|
+
|
|
282
|
+
# 1. Search child chunks
|
|
283
|
+
child_chunks = self._child_store.similarity_search( query, k=self.k )
|
|
284
|
+
# 2. Extract unique parent IDs in relevance order
|
|
285
|
+
parent_ids = list(dict.fromkeys( # standard trick to remove duplicates from a list or iterable while preserving the original item orde
|
|
286
|
+
chunk.metadata["parent_id"] for chunk in child_chunks if "parent_id" in chunk.metadata
|
|
287
|
+
))
|
|
288
|
+
if not parent_ids:
|
|
289
|
+
return []
|
|
290
|
+
# 3. Fetch parent documents
|
|
291
|
+
parents = self._parent_store.get_by_ids( parent_ids )
|
|
292
|
+
parents_by_id = { doc.id: doc for doc in parents if doc.id is not None }
|
|
293
|
+
# 4. Restore relevance order
|
|
294
|
+
return [
|
|
295
|
+
parents_by_id[pid]
|
|
296
|
+
for pid in parent_ids if pid in parents_by_id
|
|
297
|
+
]
|
|
146
298
|
|
|
147
299
|
|
|
148
300
|
#%% Hypothetical Question Retriever
|
|
149
301
|
|
|
150
|
-
def _stable_id(
|
|
151
|
-
digest = hashlib.sha256(f"{
|
|
302
|
+
def _stable_id(source_dir: str, index: int, text: str) -> str:
|
|
303
|
+
digest = hashlib.sha256(f"{source_dir}\0{index}\0{text}".encode()).hexdigest()
|
|
152
304
|
return f"chunk-{digest[:32]}"
|
|
153
305
|
|
|
154
306
|
def _clean_questions(questions: Sequence[str]) -> list[str]:
|
|
@@ -169,9 +321,6 @@ def _chroma_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
|
|
|
169
321
|
if value is not None and isinstance(value, (str, int, float, bool))
|
|
170
322
|
}
|
|
171
323
|
|
|
172
|
-
class QuestionGenerator(Protocol):
|
|
173
|
-
"""Anything capable of generating search questions for a chunk."""
|
|
174
|
-
def generate(self, document: Document) -> list[str]: ...
|
|
175
324
|
|
|
176
325
|
class _QuestionList(BaseModel):
|
|
177
326
|
questions: list[str] = Field(description="Distinct questions answered by the text")
|
|
@@ -179,7 +328,7 @@ class _QuestionList(BaseModel):
|
|
|
179
328
|
class HypotheticalQuestionGenerator:
|
|
180
329
|
"""LLM-backed question generator using provider-supported structured output."""
|
|
181
330
|
|
|
182
|
-
def __init__(self,
|
|
331
|
+
def __init__(self, llm: BaseChatModel, *, questions_per_chunk: int = 4) -> None:
|
|
183
332
|
if questions_per_chunk < 1:
|
|
184
333
|
raise ValueError("questions_per_chunk must be positive")
|
|
185
334
|
self.questions_per_chunk = questions_per_chunk
|
|
@@ -189,7 +338,7 @@ class HypotheticalQuestionGenerator:
|
|
|
189
338
|
("human", "Text:\n{text}"),
|
|
190
339
|
]
|
|
191
340
|
)
|
|
192
|
-
self._chain = prompt |
|
|
341
|
+
self._chain = prompt | llm.with_structured_output(_QuestionList)
|
|
193
342
|
|
|
194
343
|
def generate(self, document: Document) -> list[str]:
|
|
195
344
|
result = self._chain.invoke( {"count": self.questions_per_chunk, "text": document.page_content} )
|
|
@@ -203,38 +352,40 @@ class IngestionResult:
|
|
|
203
352
|
chunk_ids: tuple[str, ...]
|
|
204
353
|
|
|
205
354
|
class HypotheticalQuestionRetriever(BaseRetriever):
|
|
206
|
-
"""Search question vectors, then resolve their chunk IDs from the chunk collection."""
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
|
|
210
|
-
|
|
355
|
+
"""Search question vectors, then resolve their chunk IDs from the chunk collection."""
|
|
356
|
+
llm: Any = None
|
|
357
|
+
persist_directory: str | Path
|
|
358
|
+
embeddings: Embeddings
|
|
359
|
+
question_collection: str
|
|
360
|
+
chunk_collection: str
|
|
211
361
|
question_k: int = 12
|
|
212
362
|
chunk_k: int = 4
|
|
213
363
|
metadata_filter: dict[str, Any] | None = None
|
|
364
|
+
create_collection_if_not_exists: bool = False
|
|
365
|
+
_question_store: Chroma = PrivateAttr()
|
|
366
|
+
_chunk_collection: Any = PrivateAttr()
|
|
367
|
+
_generator: Any = PrivateAttr()
|
|
214
368
|
|
|
215
|
-
|
|
216
|
-
client = chromadb.PersistentClient(path=str(persist_directory))
|
|
217
|
-
question_store = Chroma( client=client, collection_name=question_collection, embedding_function=embeddings )
|
|
218
|
-
chunks = client.get_or_create_collection( name=chunk_collection, embedding_function=None )
|
|
219
|
-
super().__init__(
|
|
220
|
-
question_store=question_store,
|
|
221
|
-
chunk_collection=chunks,
|
|
222
|
-
generator=generator,
|
|
223
|
-
question_k=question_k,
|
|
224
|
-
chunk_k=chunk_k,
|
|
225
|
-
metadata_filter=metadata_filter,
|
|
226
|
-
**kwargs,
|
|
227
|
-
)
|
|
369
|
+
model_config = ConfigDict(arbitrary_types_allowed=True)
|
|
228
370
|
|
|
229
|
-
def
|
|
230
|
-
|
|
231
|
-
|
|
371
|
+
def __init__(self, **kwargs) -> None:
|
|
372
|
+
super().__init__(**kwargs)
|
|
373
|
+
client = chromadb.PersistentClient( path= self.persist_directory )
|
|
374
|
+
self._question_collection = Chroma( client=client, collection_name=self.question_collection, embedding_function=self.embeddings, create_collection_if_not_exists=self.create_collection_if_not_exists )
|
|
375
|
+
self._chunk_collection = client.get_or_create_collection( name=self.chunk_collection, embedding_function=None )
|
|
376
|
+
|
|
377
|
+
def ingest( self, documents: Sequence[Document], *, source_dir: str, split: bool = True, chunk_size: int = 500, chunk_overlap: int = 150, replace_source: bool = True ) -> IngestionResult:
|
|
378
|
+
if not self.llm:
|
|
379
|
+
raise ValueError("To ingest, llm must be provided in the constructor.")
|
|
380
|
+
self._generator = HypotheticalQuestionGenerator( self.llm )
|
|
381
|
+
if not source_dir: # source_dir is a unique identifier for the source document or documents being ingested
|
|
382
|
+
raise ValueError("source_dir must not be empty")
|
|
232
383
|
splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, add_start_index=True )
|
|
233
384
|
input_docs = splitter.split_documents( list(documents) ) if split else list(documents)
|
|
234
385
|
|
|
235
386
|
if replace_source:
|
|
236
|
-
self.
|
|
237
|
-
self.
|
|
387
|
+
self._question_collection.delete( where={"source_dir": source_dir} )
|
|
388
|
+
self._chunk_collection.delete( where={"source_dir": source_dir} )
|
|
238
389
|
|
|
239
390
|
chunk_ids: list[str] = []
|
|
240
391
|
question_docs: list[Document] = []
|
|
@@ -243,32 +394,31 @@ class HypotheticalQuestionRetriever(BaseRetriever):
|
|
|
243
394
|
chunk_texts: list[str] = []
|
|
244
395
|
|
|
245
396
|
for index, doc in enumerate(input_docs):
|
|
246
|
-
chunk_id = _stable_id(
|
|
247
|
-
metadata = _chroma_metadata( {**doc.metadata, "
|
|
397
|
+
chunk_id = _stable_id(source_dir, index, doc.page_content)
|
|
398
|
+
metadata = _chroma_metadata( {**doc.metadata, "source_dir": source_dir, "chunk_id": chunk_id} )
|
|
248
399
|
chunk_ids.append( chunk_id )
|
|
249
400
|
chunk_metadatas.append( metadata )
|
|
250
401
|
chunk_texts.append( doc.page_content )
|
|
251
|
-
|
|
252
|
-
questions = _clean_questions( self.generator.generate(doc) )
|
|
402
|
+
questions = _clean_questions( self._generator.generate(doc) )
|
|
253
403
|
for question_index, question in enumerate(questions):
|
|
254
404
|
question_ids.append( f"{chunk_id}:q:{question_index}" )
|
|
255
405
|
question_docs.append(
|
|
256
406
|
Document( page_content=question,
|
|
257
|
-
metadata={"chunk_id": chunk_id, "
|
|
407
|
+
metadata={"chunk_id": chunk_id, "source_dir": source_dir, "question_index": question_index },
|
|
258
408
|
)
|
|
259
409
|
)
|
|
260
410
|
if question_docs: # LangChain's Chroma.add_documents() triggers embedding
|
|
261
|
-
self.
|
|
411
|
+
self._question_collection.add_documents( question_docs, ids=question_ids )
|
|
262
412
|
|
|
263
413
|
if chunk_ids: # Chroma's underlying collection API stores text and metadata without embeddings.
|
|
264
|
-
self.
|
|
414
|
+
self._chunk_collection.upsert( ids=chunk_ids, documents=chunk_texts, metadatas=chunk_metadatas,
|
|
265
415
|
embeddings=[[0.0] for _ in chunk_ids] ) # placeholder vectors since Chroma records require vectors
|
|
266
416
|
|
|
267
417
|
return IngestionResult(len(chunk_ids), len(question_docs), tuple(chunk_ids))
|
|
268
418
|
|
|
269
419
|
# override _get_relevant_documents() method from BaseRetriever
|
|
270
420
|
def _get_relevant_documents(self, query: str, *, run_manager: Any = None) -> list[Document]:
|
|
271
|
-
matches = self.
|
|
421
|
+
matches = self._question_collection.similarity_search( query, k=self.question_k, filter=self.metadata_filter )
|
|
272
422
|
chunk_ids: list[str] = []
|
|
273
423
|
seen: set[str] = set()
|
|
274
424
|
for m in matches:
|
|
@@ -281,74 +431,16 @@ class HypotheticalQuestionRetriever(BaseRetriever):
|
|
|
281
431
|
if not chunk_ids:
|
|
282
432
|
return []
|
|
283
433
|
|
|
284
|
-
records = self.
|
|
285
|
-
|
|
434
|
+
records = self._chunk_collection.get( ids=chunk_ids, include=["documents", "metadatas"] )
|
|
435
|
+
records_by_id = {
|
|
286
436
|
record_id: Document(page_content=text, metadata=metadata or {})
|
|
287
437
|
for record_id, text, metadata in zip(
|
|
288
438
|
records["ids"], records["documents"], records["metadatas"], strict=True
|
|
289
439
|
)
|
|
290
440
|
}
|
|
291
|
-
return [
|
|
292
|
-
|
|
293
|
-
#%% RAG Chains
|
|
294
|
-
|
|
295
|
-
system_prompt = """You are a helpful assistant answering questions based on provided document snippets.
|
|
296
|
-
IMPORTANT INSTRUCTIONS:
|
|
297
|
-
1. Base your answer ONLY on the provided context snippets below.
|
|
298
|
-
2. Synthesize information from multiple snippets if needed.
|
|
299
|
-
3. If the snippets contain partial information, provide the best answer you can with what's available.
|
|
300
|
-
4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
|
|
301
|
-
5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
|
|
302
|
-
|
|
303
|
-
The snippets are provided in this format:
|
|
304
|
-
Snippet 1: <text content>
|
|
305
|
-
Snippet 2: <text content>
|
|
306
|
-
...
|
|
307
|
-
|
|
308
|
-
{context}"""
|
|
309
|
-
|
|
310
|
-
def format_documents(documents: Sequence[Document]) -> str:
|
|
311
|
-
snippets = []
|
|
312
|
-
for i, doc in enumerate(documents, start=1):
|
|
313
|
-
snippets.append(f"Snippet {i}: {doc.page_content}")
|
|
314
|
-
context = "\n".join(snippets)
|
|
315
|
-
return context
|
|
316
|
-
|
|
317
|
-
def rag_basic(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
|
|
318
|
-
"""Return an LCEL chain accepting a question string and producing an answer string."""
|
|
319
|
-
prompt = ChatPromptTemplate.from_messages(
|
|
320
|
-
[
|
|
321
|
-
( "system", "Answer using only the context below. If the answer is absent, say you do not know.\n\nContext:\n{context}" ),
|
|
322
|
-
("human", "{query}"),
|
|
323
|
-
]
|
|
324
|
-
)
|
|
325
|
-
inputs = {
|
|
326
|
-
"context": retriever | format_documents,
|
|
327
|
-
"query": RunnablePassthrough(),
|
|
328
|
-
}
|
|
329
|
-
return inputs | prompt | model | StrOutputParser()
|
|
441
|
+
return [ records_by_id[chunk_id] for chunk_id in chunk_ids if chunk_id in records_by_id ]
|
|
330
442
|
|
|
331
443
|
|
|
332
|
-
def rag_main(retriever: BaseRetriever, model: BaseChatModel) -> Runnable:
|
|
333
|
-
"""Return an LCEL chain accepting a question string and producing an dic with 4 keys/values
|
|
334
|
-
-- 'sources': original list of Documents from retriever
|
|
335
|
-
-- 'query': original query
|
|
336
|
-
-- 'context': merged content from relevant documents
|
|
337
|
-
-- 'answer': final answer """
|
|
338
|
-
prompt = ChatPromptTemplate.from_messages([
|
|
339
|
-
("system", system_prompt),
|
|
340
|
-
("human", "{query}")
|
|
341
|
-
])
|
|
342
|
-
return (
|
|
343
|
-
{
|
|
344
|
-
"sources": retriever,
|
|
345
|
-
"query": RunnablePassthrough(),
|
|
346
|
-
}
|
|
347
|
-
| RunnablePassthrough.assign( context=lambda x: format_documents(x["sources"]) )
|
|
348
|
-
| RunnablePassthrough.assign( answer=prompt | model | StrOutputParser() )
|
|
349
|
-
)
|
|
350
|
-
|
|
351
|
-
build_rag = rag_main
|
|
352
444
|
|
|
353
445
|
#%% RAG with Routing
|
|
354
446
|
|
|
@@ -368,7 +460,7 @@ Snippet 2: <text content>
|
|
|
368
460
|
|
|
369
461
|
{context}"""
|
|
370
462
|
|
|
371
|
-
def multiRagBuilder( domain_literals, router_prompt, llm, embeddings, reranker, query_llm=None, k=3):
|
|
463
|
+
def multiRagBuilder( persist_directory, domain_literals, router_prompt, llm, embeddings, reranker, query_llm=None, k=3):
|
|
372
464
|
"""
|
|
373
465
|
domain_literals: the last one is a placeholder if the router finds no domain matches the query
|
|
374
466
|
"""
|
|
@@ -444,7 +536,7 @@ def multiRagBuilder( domain_literals, router_prompt, llm, embeddings, reranker,
|
|
|
444
536
|
retrievers = {}
|
|
445
537
|
retriever_nodes = {}
|
|
446
538
|
for domain in domains:
|
|
447
|
-
retrievers[domain] = MultiQueryRerankRetriever(
|
|
539
|
+
retrievers[domain] = MultiQueryRerankRetriever( persist_directory=persist_directory, collection_name=domain, embeddings=embeddings, query_llm=query_llm, reranker=reranker, k=k)
|
|
448
540
|
retriever_nodes[domain] = partial( retriever_node, domain ) # the 1st argument is the partial function. the rest are frozen arguments passed positionally to fill the leftmost parameters of the original function.
|
|
449
541
|
builder.add_node(domain, retriever_nodes[domain])
|
|
450
542
|
builder.add_edge(domain, "generate_answer")
|
simonbb-0.0.9/pyproject.toml
DELETED
|
@@ -1,34 +0,0 @@
|
|
|
1
|
-
[build-system]
|
|
2
|
-
requires = ["hatchling >= 1.26"]
|
|
3
|
-
build-backend = "hatchling.build"
|
|
4
|
-
|
|
5
|
-
[project]
|
|
6
|
-
name = "simonbb"
|
|
7
|
-
version = "0.0.9"
|
|
8
|
-
authors = [
|
|
9
|
-
{ name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
|
|
10
|
-
]
|
|
11
|
-
description = "A small package to facilitate teaching AI courses at business schools"
|
|
12
|
-
readme = "README.md"
|
|
13
|
-
requires-python = ">=3.9"
|
|
14
|
-
classifiers = [
|
|
15
|
-
"Programming Language :: Python :: 3",
|
|
16
|
-
"Operating System :: OS Independent",
|
|
17
|
-
]
|
|
18
|
-
license = "MIT"
|
|
19
|
-
license-files = ["LICEN[CS]E*"]
|
|
20
|
-
dependencies = [
|
|
21
|
-
"matplotlib >= 3.10.8",
|
|
22
|
-
"torch >= 2.9.1",
|
|
23
|
-
"torchvision >=0.24.1",
|
|
24
|
-
"numpy >=2.4.0",
|
|
25
|
-
"transformers >= 4.57.3",
|
|
26
|
-
"gymnasium >= 1.3.0",
|
|
27
|
-
"pytorch-ignite >= 0.5.4",
|
|
28
|
-
"stable_baselines3 >= 2.9.0",
|
|
29
|
-
"ray >= 2.56.0",
|
|
30
|
-
"opencv-python >= 5.0.0.93",
|
|
31
|
-
"pyyaml >= 6.0.3"
|
|
32
|
-
]
|
|
33
|
-
[project.urls]
|
|
34
|
-
Homepage = "https://simonbb.com"
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|