An AI Powered teacher that can help you learn your topics faster before exam
import logging
from langchain.document_loaders import DirectoryLoader, PDFMinerLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceInstructEmbeddings
from langchain.vectorstores import Chroma
from config import (
    PERSIST_DIRECTORY,
    MODEL_DIRECTORY,
    SOURCE_DIR,
    EMBEDDING_MODEL,
    DEVICE_TYPE,
    CHROMA_SETTINGS,
)

def load_docs(directory: str = SOURCE_DIR):
    """
    Loads documents from a specified directory.
    
    Args:
        directory (str): The directory path containing PDF documents.
    
    Returns:
        list: A list of loaded documents.
    """
    loader = DirectoryLoader(directory, glob="**/*.pdf", use_multithreading=True, loader_cls=PDFMinerLoader)
    docs = loader.load()
    logging.info(f"Loaded {len(docs)} documents from {directory}")
    return docs

def split_docs(documents,chunk_size=1000,chunk_overlap=200):
    """
    Splits documents into smaller chunks for processing.
    
    Args:
        documents (list): List of documents to be split.
        chunk_size (int): The size of each chunk.
        chunk_overlap (int): The overlap between adjacent chunks.
    
    Returns:
        list: List of split documents.
    """
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    docs = text_splitter.split_documents(documents)
    logging.info(f"Split {len(docs)} documents into chunks")
    return docs

def builder():
    """
    Builds the database by loading, splitting, and embedding documents.
    """
    logging.info("Building the database")
    documents = load_docs()
    docs = split_docs(documents)

    embeddings = HuggingFaceInstructEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"device": DEVICE_TYPE},
        cache_folder=MODEL_DIRECTORY,
    )
    db = Chroma.from_documents(
        docs,
        embeddings,
        persist_directory=PERSIST_DIRECTORY,
        client_settings=CHROMA_SETTINGS,

    )
    logging.info(f"Loaded Documents to Chroma DB Successfully")


if __name__ == '__main__':
 
    builder()
builder.py
Priyamakeshwari/TeachGPT
import os
import torch
from chromadb.config import Settings

# Source Directory for Documents to Ingest
SOURCE_DIR = os.path.join(os.path.dirname(__file__), "documents")
# To store models from HuggingFace
MODEL_DIRECTORY = os.path.join(os.path.dirname(__file__), "models")
# Store Vector Embeddings in db directory
PERSIST_DIRECTORY = os.path.join(os.path.dirname(__file__), "db")


# MODELS 
MODEL_NAME = "TheBloke/Mistral-7B-Instruct-v0.1-GGUF"
MODEL_FILE = "mistral-7b-instruct-v0.1.Q4_K_M.gguf"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L12-v2"

# MODEL CONFIG
MAX_TOKEN_LENGTH = 4094 # 8192 is the max for Mistral-7B
N_GPU_LAYERS = 40

# PYTORCH DEVICE COMPATIBILITY
if torch.cuda.is_available():
    DEVICE_TYPE = "cuda"
elif torch.backends.mps.is_available():
    DEVICE_TYPE = "mps"
else:
    DEVICE_TYPE = "cpu"

# DATABASE SETTINGS
CHROMA_SETTINGS = Settings(
    anonymized_telemetry=False,
    is_persistent=True,
)

config.py
Priyamakeshwari/TeachGPT
import logging
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.chains import RetrievalQA, RetrievalQAWithSourcesChain
from langchain.llms import LlamaCpp
from langchain.embeddings import HuggingFaceInstructEmbeddings
from langchain.callbacks.manager import CallbackManager
from langchain.vectorstores import Chroma
from langchain.prompts import PromptTemplate
from huggingface_hub import hf_hub_download

from config import (
    PERSIST_DIRECTORY,
    MODEL_DIRECTORY,
    SOURCE_DIR,
    EMBEDDING_MODEL,
    DEVICE_TYPE,
    CHROMA_SETTINGS,
    MODEL_NAME,
    MODEL_FILE,
    N_GPU_LAYERS,
    MAX_TOKEN_LENGTH,
)

def load_model(device_type:str = DEVICE_TYPE, model_id:str = MODEL_NAME, model_basename:str = MODEL_FILE, LOGGING=logging):
    """
    Load a language model.

    Args:
        device_type (str): The type of device to use ('cuda', 'mps', or 'cpu').
        model_id (str): The ID of the model to load.
        model_basename (str): The name of the model file.
        LOGGING (logging): The logging object.

    Returns:
        LlamaCpp: The loaded language model.
    """
    callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
    try:
        model_path = hf_hub_download(
            repo_id=model_id,
            filename=model_basename,
            resume_download=True,
            cache_dir=MODEL_DIRECTORY,
        )
        kwargs = {
            "model_path": model_path,
            "max_tokens": MAX_TOKEN_LENGTH,
            "n_ctx": MAX_TOKEN_LENGTH,
            "n_batch": 512,  
            "callback_manager": callback_manager,
            "verbose":False,
            "f16_kv":True,
            "streaming":True,
        }
        if device_type.lower() == "mps":
            kwargs["n_gpu_layers"] = 1
        if device_type.lower() == "cuda":
            kwargs["n_gpu_layers"] = N_GPU_LAYERS  # set this based on your GPU
        llm =  LlamaCpp(**kwargs)
        LOGGING.info(f"Loaded {model_id} locally")
        return llm  # Returns a LlamaCpp object
    except Exception as e:
        LOGGING.info(f"Error {e}")

def retriver(device_type:str = DEVICE_TYPE, LOGGING=logging):
    """
    Retrieve information using a language model and Chroma database.

    Args:
        device_type (str): The type of device to use ('cuda', 'mps', or 'cpu').
        LOGGING (logging): The logging object.
    """
    embeddings = HuggingFaceInstructEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"device": DEVICE_TYPE},
        cache_folder=MODEL_DIRECTORY,
    )
    db = Chroma(
        persist_directory=PERSIST_DIRECTORY,
        embedding_function=embeddings,
    )
    retriever = db.as_retriever()
    LOGGING.info(f"Loaded Chroma DB Successfully")
    llm = load_model(device_type, model_id=MODEL_NAME, model_basename=MODEL_FILE, LOGGING=logging)
    template = """
    [INST]
        Context: {summaries}
        User: {question}
    [/INST]
    """ 
    prompt = PromptTemplate(input_variables=["summaries", "question"], template=template)
    chain = RetrievalQAWithSourcesChain.from_chain_type(
        llm=llm,
        retriever=retriever,
        # chain_type="stuff",
        chain_type_kwargs={"prompt": prompt},
    )

    chain({'question'  : "What is the linux command to list files in direcotyu",},return_only_outputs=True)


if __name__ == '__main__':
    retriver()

    

main.py
Priyamakeshwari/TeachGPT