[MLC-14] server: step toward Retrieval Augmented Generation (RAG) w/ …

[stockeh]

…indexing (load, split, store) retriever.

Overview of RAG: https://python.langchain.com/docs/use_cases/question_answering/

Implementation Details

• Extends langchain with mlx adaptions: RecursiveCharacterTextSplitter, OpenAIEmbeddings, Chroma
• Custom Embeddings from Gemma model, averaged vector chunk tokens to have unit length (ideal for cosine similarity)

Risks

• Similarity score: empirical favoring toward max_marginal_relevance_search over similarity_search
• Gemma embedding layer isn't optimized for text retrieval. Should we use an alternative text embedding model for this?

[stockeh]

Example usage of retriever:

from server.utils import load

from server.retriever.loader import directory_loader
from server.retriever.splitter import RecursiveCharacterTextSplitter
from server.retriever.vectorstore import Chroma, Embeddings


def main():
    model, tokenizer = load('mlx-community/quantized-gemma-7b-it')
    raw_docs = directory_loader(
        '/Users/stock/Library/Mobile Documents/iCloud~md~obsidian/Documents/main')
    print(len(raw_docs), len(raw_docs[0].page_content))
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1024, chunk_overlap=32, add_start_index=True
    )
    splits = text_splitter.split_documents(raw_docs)
    print(len(splits), len(splits[0].page_content), splits[0].metadata)
    db = Chroma.from_documents(
        documents=splits, embedding=Embeddings(model.model, tokenizer))
    print('-------------------')
    query = "What is a cascade neural network?"
    # docs = db.similarity_search(query)
    docs = db.max_marginal_relevance_search(query)
    print('>', query)
    for doc in docs:
        print(doc.page_content, doc.metadata, sep='\n')
        print('-------------------')


if __name__ == '__main__':
    main()