Retrieval Augmented Generation of Text Using Embeddings
Retrieval-Augmented Generation (RAG) is a framework that combines the strengths of pre-trained language models (LLMs) with retrievers. Retrievers are system components for accessing knowledge from external sources of text data. In RAG a retriever selects relevant documents or passages from a corpus, and a generator produces a response based on both the retrieved information and the input query. The process typically follows these steps that we will use in the example Racket code:
- Query Encoding: The input query is encoded into a vector representation.
- Document Retrieval: A retriever system uses the query representation to fetch relevant documents or passages from an external corpus.
- Document Encoding: The retrieved documents are encoded into vector representations.
- Joint Encoding: The query and document representations are combined, often concatenated or mixed via attention mechanisms.
- Generation: A generator, usually LLM, is used to produce a response based on the joint representation.
RAG enables the LLM to access and leverage external text data sources, which is crucial for tasks that require information beyond what the LLM has been trained on. It’s a blend of retrieval-based and generation-based approaches, aimed at boosting the factual accuracy and informativeness of generated responses.
Example Implementation
In the following short Racket example program (file Racket-AI-book/source-code/embeddingsdb /embeddingsdb.rkt) I implement some ideas of a RAG architecture. At file load time the text files in the subdirectory data are read, split into “chunks”, and each chunk along with its parent file name and OpenAI text embedding is stored in a local SQLite database. When a user enters a query, the OpenAI embedding is calculated, and this embedding is matched against the embeddings of all chunks using the dot product of two 1536 element embedding vectors. The “best” chunks are concatenated together and this “context” text is passed to GPT-4 along with the user’s original query. Here I describe the code in more detail:
The provided Racket code uses a local SQLite database and OpenAI’s APIs for calculating text embeddings and for text completions.
Utility Functions:
floats->stringandstring->floatsare utility functions for converting between a list of floats and its string representation.read-filereads a file’s content.join-stringsjoins a list of strings with a specified separator.truncate-stringtruncates a string to a specified length.interleavemerges two lists by interleaving their elements.break-into-chunksbreaks a text into chunks of a specified size.string-to-listanddecode-roware utility functions for parsing and processing database rows.
Database Setup:
- Database connection is established to “test.db” and a table named “documents” is created with columns for document_path, content, and embedding.
Document Management:
insert-documentinserts a document and its associated information into the database.get-document-by-document-pathandall-documentsare utility functions for querying documents from the database.create-documentreads a document from a file path, breaks it into chunks, computes embeddings for each chunk via a functionembeddings-openai, and inserts these into the database.
Semantic Matching and Interaction:
execute-to-listanddot-productare utility functions for database queries and vector operations.semantic-matchperforms a semantic search by calculating the dot product of embeddings of the query and documents in the database. It then aggregates contexts of documents with a similarity score above a certain threshold, and sends a new query constructed with these contexts to OpenAI for further processing.QAis a wrapper aroundsemantic-matchfor querying.CHATinitiates a loop for user interaction where each user input is processed throughsemantic-matchto generate a response, maintaining a context of the previous chat.
Test Code:
testfunction creates documents by reading from specified file paths, and performs some queries using theQAfunction.
The code uses a local SQLite database to store and manage document embeddings and the OpenAI API for generating embeddings and performing semantic searches based on user queries. Two functions are exported in case you want to use this example as a library: create-document and QA.
1 #lang racket
2
3 (require db)
4 (require llmapis)
5 (require racket/runtime-path)
6
7 (provide create-document QA CHAT semantic-match)
8
9 ; Function to convert list of floats to string representation
10 (define (floats->string floats)
11 (string-join (map number->string floats) " "))
12
13 ; Function to convert string representation back to list of floats
14 (define (string->floats str)
15 (map string->number (string-split str)))
16
17 (define (read-file infile)
18 (with-input-from-file infile
19 (lambda ()
20 (let ((contents (read)))
21 contents))))
22
23 (define (join-strings separator list)
24 (string-join list separator))
25
26 (define (truncate-string string length)
27 (substring string 0 (min length (string-length string))))
28
29 (define (interleave list1 list2)
30 (if (or (null? list1) (null? list2))
31 (append list1 list2)
32 (cons (car list1)
33 (cons (car list2)
34 (interleave (cdr list1) (cdr list2))))))
35
36 (define (break-into-chunks text chunk-size)
37 (let loop ((start 0) (chunks '()))
38 (if (>= start (string-length text))
39 (reverse chunks)
40 (loop (+ start chunk-size)
41 (cons (substring text start (min (+ start chunk-size) (string-length text))) chunks)))))
42
43 (define (string-to-list str)
44 (map string->number (string-split str)))
45
46 (define (decode-row row)
47 (let ((id (vector-ref row 0))
48 (context (vector-ref row 1))
49 (embedding (string-to-list (vector-ref row 2))))
50 (list id context embedding)))
51
52 (define db (sqlite3-connect #:database "test.db" #:mode 'create #:use-place #t))
53
54 (with-handlers ([exn:fail? (lambda (ex) (void))])
55 (query-exec
56 db
57 "CREATE TABLE documents (document_path TEXT, content TEXT, embedding TEXT);"))
58
59 (define (insert-document document-path content embedding)
60 (printf "~%insert-document:~% content:~a~%~%" content)
61 (query-exec
62 db
63 "INSERT INTO documents (document_path, content, embedding) VALUES (?, ?, ?);"
64 document-path content (floats->string embedding)))
65
66 (define (get-document-by-document-path document-path)
67 (map decode-row
68 (query-rows db
69 "SELECT * FROM documents WHERE document_path = ?;"
70 document-path)))
71
72 (define (all-documents)
73 (map
74 decode-row
75 (query-rows
76 db
77 "SELECT * FROM documents;")))
78
79 (define (create-document fpath)
80 (let ((contents (break-into-chunks (file->string fpath) 200)))
81 (for-each
82 (lambda (content)
83 (with-handlers ([exn:fail? (lambda (ex) (void))])
84 (let ((embedding (embeddings-openai content)))
85 (insert-document fpath content embedding))))
86 contents)))
87
88 ;; Assuming a function to fetch documents from database
89 (define (execute-to-list db query)
90 (query-rows db query))
91
92 (define (dot-product a b) ;; dot product of two lists of floating point numbers
93 (for/sum ([x a] [y b])
94 (* x y)))
95
96 (define (semantic-match query custom-context [cutoff 0.7])
97 (let ((emb (embeddings-openai query))
98 (ret '()))
99 (for-each
100 (lambda (doc)
101 (let* ((context (second doc))
102 (embedding (third doc))
103 (score (dot-product emb embedding)))
104 (when (> score cutoff)
105 (set! ret (cons context ret)))))
106 (all-documents))
107 (printf "~%semantic-search: ret=~a~%" ret)
108 (let* ((context (string-join (reverse ret) " . "))
109 (query-with-context (string-join (list context custom-context "Question:" query) " ")))
110 (question-openai query-with-context))))
111
112 (define (QA query [quiet #f])
113 (let ((answer (semantic-match query "")))
114 (unless quiet
115 (printf "~%~%** query: ~a~%** answer: ~a~%~%" query answer))
116 answer))
117
118 (define (CHAT)
119 (let ((messages '(""))
120 (responses '("")))
121 (let loop ()
122 (printf "~%Enter chat (STOP or empty line to stop) >> ")
123 (let ((string (read-line)))
124 (cond
125 ((or (string=? string "STOP") (< (string-length string) 1))
126 (list (reverse messages) (reverse responses)))
127 (else
128 (let* ((custom-context
129 (string-append
130 "PREVIOUS CHAT: "
131 (string-join (reverse messages) " ")))
132 (response (semantic-match string custom-context)))
133 (set! messages (cons string messages))
134 (set! responses (cons response responses))
135 (printf "~%Response: ~a~%" response)
136 (loop))))))))
137
138 (define-runtime-path data-dir "data")
139
140 (define (test)
141 "Test code for Semantic Document Search Using OpenAI GPT APIs and local vector database"
142 (create-document (path->string (simplify-path (build-path data-dir "sports.txt"))))
143 (create-document (path->string (simplify-path (build-path data-dir "chemistry.txt"))))
144 (QA "What is the history of the science of chemistry?")
145 (QA "What are the advantages of engaging in sports?"))
Let’s look at a few examples form a Racket REPL:
1 > (QA "What is the history of the science of chemistry?")
2 ** query: What is the history of the science of chemistry?
3 ** answer: The history of the science of chemistry dates back thousands of years. Ancient civilizations such as the Egyptians, Greeks, and Chinese were experimenting with various substances and observing chemical reactions even before the term "chemistry" was coined.
4
5 The foundations of modern chemistry can be traced back to the works of famous scholars such as alchemists in the Middle Ages. Alchemists sought to transform common metals into gold and discover elixirs of eternal life. Although their practices were often based on mysticism and folklore, it laid the groundwork for the understanding of chemical processes and experimentation.
6
7 In the 17th and 18th centuries, significant advancements were made in the field of chemistry. Prominent figures like Robert Boyle and Antoine Lavoisier began to understand the fundamental principles of chemical reactions and the concept of elements. Lavoisier is often referred to as the "father of modern chemistry" for his work in establishing the law of conservation of mass and naming and categorizing elements.
8
9 Throughout the 19th and 20th centuries, chemistry continued to progress rapidly. The development of the periodic table by Dmitri Mendeleev in 1869 revolutionized the organization of elements. The discovery of new elements, the formulation of atomic theory, and the understanding of chemical bonding further expanded our knowledge.
10
11 Chemistry also played a crucial role in various industries and technologies, such as the development of synthetic dyes, pharmaceuticals, plastics, and materials. The emergence of quantum mechanics and spectroscopy in the early 20th century opened up new avenues for understanding the behavior of atoms and molecules.
12
13 Today, chemistry is an interdisciplinary science that encompasses various fields such as organic chemistry, inorganic chemistry, physical chemistry, analytical chemistry, and biochemistry. It continues to evolve and make significant contributions to society, from developing sustainable materials to understanding biological processes and addressing global challenges such as climate change.
14
15 In summary, the history of the science of chemistry spans centuries, starting from ancient civilizations to the present day, with numerous discoveries and advancements shaping our understanding of the composition, properties, and transformations of matter.
This output is the combination of data found in the text files in the directory Racket-AI-book/source-code/embeddingsdb/data and the data that OpenAI GPT-4 was trained on. Since the local “document” file chemistry.txt is very short, most of this output is derived from the innate knowledge GPT-4 has from its training data.
In order to show that this example is also using data in the local “document” text files, I manually edited the file data/chemistry.txt adding the following made-up organic compound:
1 ZorroOnian Alcohol is another organic compound with the formula C 6 H 10 O.
GPT-4 was never trained on my made-up data so it has no idea what the non-existent compound ZorroOnian Alcohol is. The following answer is retrieved via RAG from the local document data (for brevity, most of the output for adding the local document files to the embedding index is not shown):
1 > (create-document
2 "/Users/markw/GITHUB/Racket-AI-book/source-code/embeddingsdb/data/chemistry.txt")
3
4 insert-document:
5 content:Amyl alcohol is an organic compound with the formula C 5 H 12 O. ZorroOnian Alcohol is another organic compound with the formula C 6 H 10 O. All eight isomers of amyl alcohol are known.
6
7 ...
8
9 > (QA "what is the formula for ZorroOnian Alcohol")
10
11 ** query: what is the formula for ZorroOnian Alcohol
12 ** answer: The formula for ZorroOnian Alcohol is C6H10O.
There is also a chat interface:
1 Enter chat (STOP or empty line to stop) >> who is the chemist Robert Boyle
2
3 Response: Robert Boyle was an Irish chemist and physicist who is known as one of the pioneers of modern chemistry. He is famous for Boyle's Law, which describes the inverse relationship between the pressure and volume of a gas, and for his experiments on the properties of gases. He lived from 1627 to 1691.
4
5 Enter chat (STOP or empty line to stop) >> Where was he born?
6
7 Response: Robert Boyle was born in Lismore Castle, County Waterford, Ireland.
8
9 Enter chat (STOP or empty line to stop) >>
The following diagram shows the high-level architecture of the RAG pipeline developed in this chapter:
Retrieval Augmented Generation Wrap Up
Retrieval Augmented Generation (RAG) is one of the best use cases for semantic search. Another way to write RAG applications is to use a web search API to get context text for a query, and add this context data to whatever context data you have in a local embeddings data store.
Optional Practice Problems
- Calculate Cosine Similarity: In
embeddingsdb.rkt, semantic similarity is measured using the rawdot-product. Implement a truecosine-similaritymetric that divides the dot product by the product of the vector magnitudes, ensuring compatibility with embeddings that are not pre-normalized to unit length. - Chunking with Overlaps: The current
break-into-chunksimplementation slices text at arbitrary character boundaries. Write an improved chunking function that slices text at sentence or paragraph boundaries and allows for a configurable overlap (e.g., 200 characters with a 40-character overlap) to preserve local context. - Extend Database Operations: Add a function
delete-documentthat deletes all chunks and vector representations associated with a given file path from the SQLite database.