huberLLM - Retrieval Augmented LLM for Huberman Podcast

import os
import torch

from langchain.llms.base import LLM
from transformers import T5Tokenizer, T5ForConditionalGeneration, GPT2Tokenizer, GPT2LMHeadModel, AutoTokenizer, AutoModelForCausalLM

import langchain 
langchain.__version__

from transformers import pipeline

"""
if filename.endswith(".txt"):
        with open(os.path.join(source_directory, filename), 'r') as f:
            s = f.read()
            with open(os.path.join(target_directory, filename), "w") as text_file:
                text_file.write(s.replace("{", "").replace("{", "").replace("[", "").replace("]", ""))
"""
root = "data/"
source_dir = "transcripts/"
target_dir = "cleaned_transcripts/"

source_path = os.path.join(root, source_dir)
target_path = os.path.join(root, target_dir)

os.makedirs(target_path, exist_ok=True)

with os.scandir(os.path.join(root, source_dir)) as it:
    for entry in it:
        if entry.name.endswith(".txt") and entry.is_file():
            with open(entry.path, 'r') as f:
                s = f.read()
                with open(os.path.join(target_path, entry.name), 'w') as text_file:
                    text_file.write(s.replace("{", "").replace("{", "").replace("[", "").replace("]", ""))

from langchain.document_loaders import DirectoryLoader, TextLoader

loader  = DirectoryLoader(target_path, glob="*.txt", 
                          show_progress = True, 
                          loader_cls=TextLoader) # glob parameter controls what type of files to load

docs = loader.load()

docs[0].metadata['source']

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 2048,
    chunk_overlap  = 100,
    length_function = len,
)

texts = text_splitter.create_documents(texts = [doc.page_content for doc in docs],
                                       metadatas= [doc.metadata for doc in docs] )

print(f'# documents before splitting - {len(docs)}')
print(f'# documents after splitting - {len(texts)}')

texts[0].page_content

from langchain.embeddings import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2", cache_folder='/home/ubuntu/models/')

doc_embeddings = embeddings.embed_documents(texts=[text.page_content for text in texts])

doc_embeddings[0]

from langchain.vectorstores import Chroma

db = Chroma.from_documents(documents=texts,
                           embedding=embeddings,
                           collection_name="huberdb", 
                           persist_directory="data/huberdb/")

# example query
query = "Who is lex friedman?"
docs = db.similarity_search(query)

docs[0].metadata

from langchain import HuggingFacePipeline
from transformers import T5Tokenizer, T5ForConditionalGeneration, pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM

class customLLM(LLM):
    model_name = "psmathur/orca_mini_7b"
    # tokenizer = T5Tokenizer.from_pretrained(model_name)
    # model = T5ForConditionalGeneration.from_pretrained(model_name, 
    #                                                    device_map="auto",
    #                                                    cache_dir="/home/ubuntu/models/")
    pipeline = pipeline("text-generation", 
                        model=model_name, 
                        device_map = 'auto', 
                        model_kwargs={"torch_dtype":torch.bfloat16, 
                                       "cache_dir": "/home/ubuntu/models/hugging_face/"})

    def _call(self, prompt, stop=None):
        return self.pipeline(prompt, max_length=9999)[0]["generated_text"]
 
    @property
    def _identifying_params(self):
        return {"name_of_model": self.model_name}
    @property
    def _llm_type(self):
        return "custom"

llm = customLLM()

from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.chains import RetrievalQA

# qa = ConversationalRetrievalChain.from_llm(llm, db.as_retriever(), memory=memory)

retriever_qa = RetrievalQA.from_chain_type(llm, 
                                           chain_type="stuff",
                                           retriever=db.as_retriever())

query = "Who is the main speaker in the conversation?"
result = retriever_qa(query)
print(result["result"].split('Question')[-1])

query = "Who is andrew huberman and what does he do?"
result = retriever_qa(query)
print(result["result"].split('Question')[-1])

query = "What is dopamine? Explain in detail."
result = retriever_qa(query)
print(result["result"].split('Question')[-1])