simonbb 0.0.1__tar.gz → 0.0.2__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.
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: simonbb
3
- Version: 0.0.1
4
- Summary: A small package for teaching deep learning
3
+ Version: 0.0.2
4
+ Summary: A small package for teaching AI/ML courses
5
5
  Project-URL: Homepage, https://simonbb.com
6
6
  Project-URL: Issues, https://simonbb.com/issues
7
7
  Author-email: Huaxia Rui <huaxia.rui@simon.rochester.edu>
@@ -4,11 +4,11 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "simonbb"
7
- version = "0.0.1"
7
+ version = "0.0.2"
8
8
  authors = [
9
9
  { name="Huaxia Rui", email="huaxia.rui@simon.rochester.edu" },
10
10
  ]
11
- description = "A small package for teaching deep learning"
11
+ description = "A small package for teaching AI/ML courses"
12
12
  readme = "README.md"
13
13
  requires-python = ">=3.9"
14
14
  classifiers = [
File without changes
@@ -481,6 +481,4 @@ if __name__ == "__main__":
481
481
  print(encoded)
482
482
  print(vectorizer.vocab)
483
483
  print(vectorizer.word_to_idx)
484
- print(vectorizer.idx_to_word)
485
- else:
486
- print(f'Importing from local file "{__name__}.py"')
484
+ print(vectorizer.idx_to_word)
File without changes
@@ -0,0 +1,259 @@
1
+ import os
2
+ import re
3
+ from typing import Annotated, Literal, TypedDict, get_args
4
+
5
+ from functools import partial
6
+
7
+ import torch
8
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9
+
10
+ from langchain_core.documents import Document
11
+ from langchain_core.messages import HumanMessage, SystemMessage
12
+ from langchain_core.prompts import ChatPromptTemplate
13
+ from langchain_core.runnables import RunnableLambda
14
+
15
+ from langgraph.graph import END, START, StateGraph
16
+ from langgraph.graph.message import add_messages
17
+
18
+ from langchain_chroma import Chroma
19
+ from langchain_huggingface import HuggingFaceEmbeddings
20
+
21
+ from sentence_transformers import CrossEncoder
22
+
23
+ vstore_path = f'{os.getenv("HOME")}/vstore'
24
+ embedding_model = HuggingFaceEmbeddings(
25
+ model_name="BAAI/bge-large-en-v1.5",
26
+ model_kwargs={'device': device},
27
+ encode_kwargs={'normalize_embeddings': False}
28
+ )
29
+
30
+ #%%
31
+
32
+ def load_documents(path):
33
+ """Load documents from a file or directory."""
34
+ from langchain_community.document_loaders import PyPDFLoader, UnstructuredHTMLLoader
35
+
36
+ docs = []
37
+ success_cnt = 0
38
+ page_cnt = 0
39
+
40
+ if os.path.isfile(path):
41
+ filename = os.path.basename(path)
42
+ print(filename)
43
+ ext = filename.lower()
44
+ if ext.endswith('.pdf'):
45
+ documents = PyPDFLoader(path).load()
46
+ elif ext.endswith(('.html', '.ejs')):
47
+ documents = UnstructuredHTMLLoader(path).load()
48
+ else:
49
+ print(f"Unsupported file type: {filename}")
50
+ return None, 0, 0
51
+ success_cnt = 1
52
+ page_cnt = len(documents)
53
+ docs.extend(documents)
54
+ elif os.path.isdir(path):
55
+ for filename in os.listdir(path):
56
+ file = os.path.join(path, filename)
57
+ if not os.path.isfile(file):
58
+ continue
59
+ print(filename)
60
+ ext = filename.lower()
61
+ if ext.endswith('.pdf'):
62
+ documents = PyPDFLoader(file).load()
63
+ elif ext.endswith(('.html', '.ejs')):
64
+ documents = UnstructuredHTMLLoader(file).load()
65
+ else:
66
+ continue
67
+ success_cnt += 1
68
+ page_cnt += len(documents)
69
+ docs.extend(documents)
70
+ else:
71
+ print(f"Error: '{path}' is not a valid file or directory")
72
+ return None, 0, 0
73
+
74
+ if success_cnt == 0:
75
+ print("No valid documents found!")
76
+ return None, 0, 0
77
+
78
+ return docs, success_cnt, page_cnt
79
+
80
+ reranker_model = CrossEncoder("BAAI/bge-reranker-large")
81
+ def _getRetriever( knowledge_base, query_llm=None, k=3 ):
82
+ store = Chroma( persist_directory=f'{vstore_path}/{knowledge_base}/BAAI_db/chroma_vector_store', embedding_function=embedding_model )
83
+ base_retriever = store.as_retriever(search_kwargs={"k": k}) # a vector store becomes a retriever through the as_retriever() method
84
+ if query_llm==None:
85
+ return base_retriever
86
+
87
+ perspectives_prompt = ChatPromptTemplate.from_template(
88
+ """You are an AI language model assistant. Your task is to generate 5 different versions of the given original question
89
+ These should be in 5 lines and provide ONLY these alternative questions separated by newlines, no explaination is needed.
90
+ Original question: {question}""")
91
+ def parse_queries_output(message):
92
+ #print( message.content )
93
+ return message.content.split('\n')
94
+
95
+ query_gen = perspectives_prompt | query_llm | parse_queries_output
96
+
97
+ def get_unique_union(document_lists): # Flatten list of lists, and deduplicate them
98
+ deduped_docs = { doc.page_content: doc for sublist in document_lists for doc in sublist}
99
+ unique_docs = list(deduped_docs.values()) # a flat list of unique docs
100
+ return unique_docs
101
+
102
+ def rerank_docs(query, docs, top_n=3):
103
+ """Rerank documents using cross-encoder and return top_n."""
104
+ if not docs:
105
+ return []
106
+ pairs = [(query, doc.page_content) for doc in docs]
107
+ scores = reranker_model.predict(pairs)
108
+ scored_docs = sorted(zip(scores, docs), key=lambda x: float(x[0]), reverse=True)
109
+ return [doc for _, doc in scored_docs[:top_n]]
110
+
111
+ def retrieve_and_rerank(query):
112
+ # Generate alternative queries + include original
113
+ alt_queries = query_gen.invoke({"question": query})
114
+ all_queries = [query] + alt_queries
115
+ # Retrieve chunks for each query
116
+ doc_lists = base_retriever.batch(all_queries)
117
+ unique_docs = get_unique_union(doc_lists)
118
+ # Rerank with cross-encoder and return top 3
119
+ return rerank_docs(query, unique_docs, top_n=5)
120
+
121
+ return RunnableLambda(retrieve_and_rerank)
122
+
123
+
124
+ """
125
+ RAG Chain
126
+ """
127
+
128
+ def getRagModel(knowledge_base, llm, query_llm=None, k=3):
129
+ """
130
+ Create and return a RAG model that returns a dictionary with 'result' and 'source_documents'.
131
+ """
132
+ retriever = _getRetriever(knowledge_base, query_llm, k)
133
+
134
+ def _rag_chain(input_dict):
135
+ query = input_dict["query"]
136
+ docs = retriever.invoke(query)
137
+ snippets = []
138
+ for i, doc in enumerate(docs, start=1):
139
+ snippets.append(f"Snippet {i}: {doc.page_content}")
140
+ context = "\n".join(snippets)
141
+ system_prompt = """You are a helpful assistant answering questions based on provided document snippets.
142
+ IMPORTANT INSTRUCTIONS:
143
+ 1. Base your answer ONLY on the provided context snippets below.
144
+ 2. Synthesize information from multiple snippets if needed.
145
+ 3. If the snippets contain partial information, provide the best answer you can with what's available.
146
+ 4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
147
+ 5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
148
+
149
+ The snippets are provided in this format:
150
+ Snippet 1: <text content>
151
+ Snippet 2: <text content>
152
+ ...
153
+
154
+ {context}"""
155
+ prompt = ChatPromptTemplate.from_messages( [ ("system", system_prompt), ("human", "{query}") ])
156
+ messages = prompt.format_messages( context=context, query=query )
157
+ answer = llm.invoke(messages).content
158
+ return {"answer": answer, "sources": docs}
159
+
160
+ return RunnableLambda(_rag_chain)
161
+
162
+
163
+ """
164
+ multi-rag system with router
165
+ """
166
+
167
+ domain_prompt = """You are a helpful assistant answering questions based on provided document snippets relating to {domain}.
168
+
169
+ IMPORTANT INSTRUCTIONS:
170
+ 1. Base your answer ONLY on the provided context snippets below.
171
+ 2. Synthesize information from multiple snippets if needed.
172
+ 3. If the snippets contain partial information, provide the best answer you can with what's available.
173
+ 4. Only say "There wasn't enough information to answer the question" if the snippets truly contain NO relevant information.
174
+ 5. Always cite your sources using the snippet numbers in parentheses, like (1), (2), etc.
175
+
176
+ The snippets are provided in this format:
177
+ Snippet 1: <text content>
178
+ Snippet 2: <text content>
179
+ ...
180
+
181
+ {context}"""
182
+
183
+ def multiRagBuilder( domain_literals, router_prompt, chat_router, chat, query_llm=None, k=3):
184
+ """
185
+ domain_literals: the last one is a placeholder if the router finds no domain matches the query
186
+ """
187
+ domains = get_args(domain_literals)[:-1] # exclude the placeholder
188
+ nodomain = get_args(domain_literals)[-1]
189
+
190
+ class State(TypedDict):
191
+ messages: Annotated[list, add_messages] # annotated with the reducer function to append new messages to the existing list, rather than overwriting
192
+ query: str # input
193
+ domain: domain_literals # output
194
+ documents: list[Document]
195
+ answer: str
196
+
197
+ class Input(TypedDict):
198
+ query: str
199
+
200
+ class Output(TypedDict):
201
+ domain: domain_literals
202
+ documents: list[Document]
203
+ answer: str
204
+
205
+ def router_node(state: State) -> State:
206
+ user_message = HumanMessage(state["query"])
207
+ messages = [ SystemMessage(router_prompt), *state["messages"], user_message ]
208
+ res = chat_router.invoke(messages)
209
+ return {
210
+ "domain": res.content.strip(),
211
+ "messages": [user_message, res], # update conversation history
212
+ }
213
+
214
+ def pick_retriever( state: State) -> Literal[*domains, "generate_answer"]:
215
+ for domain in domains:
216
+ if re.search( domain, state['domain'] ):
217
+ return domain
218
+ return 'generate_answer'
219
+
220
+ def generate_answer(state: State) -> State:
221
+ snippets = []
222
+ for i, doc in enumerate(state.get('documents', []), start=1):
223
+ snippets.append(f"Snippet {i}: {doc.page_content}")
224
+ context = "\n".join(snippets)
225
+ if state.get('domain', '')!=nodomain:
226
+ prompt = ChatPromptTemplate.from_messages( [ ("system", domain_prompt), ("human", "{query}") ])
227
+ messages = prompt.format_messages( domain=state.get('domain', ''), context=context, query=state["query"] )
228
+ else:
229
+ prompt = ChatPromptTemplate.from_messages( [ ("system", 'You are a helpful assistant.'), ("human", "{query}") ])
230
+ messages = prompt.format_messages( query=state["query"] )
231
+ res = chat.invoke(messages)
232
+ return {
233
+ "answer": res.content,
234
+ "messages": res, # update conversation history
235
+ }
236
+
237
+ def retriever_node( domain: domains, state: State):
238
+ query = re.sub(domain, 'this class', state["query"] ) # using the course code hurts RAG since the course code is rarely used in documents.
239
+ #print( f"domain={domain}; query={query}")
240
+ documents = retrievers[domain].invoke( query )
241
+ return { "documents": documents, 'query': query }
242
+
243
+ builder = StateGraph(State, input_schema=Input, output_schema=Output)
244
+ builder.add_node("router", router_node)
245
+ builder.add_edge(START, "router")
246
+ builder.add_conditional_edges("router", pick_retriever)
247
+
248
+ retrievers = {}
249
+ retriever_nodes = {}
250
+ for domain in domains:
251
+ retrievers[domain] = _getRetriever(domain, query_llm, k)
252
+ 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.
253
+ builder.add_node(domain, retriever_nodes[domain])
254
+ builder.add_edge(domain, "generate_answer")
255
+
256
+ builder.add_node("generate_answer", generate_answer)
257
+ builder.add_edge("generate_answer", END)
258
+
259
+ return builder
File without changes
File without changes
File without changes