gpt.py

# -*- coding: utf-8 -*-
"""GPT.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Is95K1fipQXeOQ-5S6JI7moTVRA9jW0-
"""

!pip install -q git+https://github.com/huggingface/transformers
!pip install -qU langchain Faiss-gpu tiktoken sentence-transformers
!pip install -qU trl Py7zr auto-gptq optimum
!pip install -q rank_bm25
!pip install -q PyPdf

from google.colab import drive

drive.mount('/content/gdrive/', force_remount=True)

!pip install pymupdf
!pip install gradio

from langchain.embeddings import CacheBackedEmbeddings,HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.storage import LocalFileStore
from langchain.retrievers import BM25Retriever,EnsembleRetriever
from langchain.document_loaders import PyMuPDFLoader,DirectoryLoader
from langchain.llms import HuggingFacePipeline
from langchain.cache import InMemoryCache
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import prompt
from langchain.chains import RetrievalQA
from langchain.callbacks import StdOutCallbackHandler
from langchain import PromptTemplate
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

dir_loader = DirectoryLoader("/content/data",
                             glob="*.pdf",
                             loader_cls=PyMuPDFLoader)
docs = dir_loader.load()

text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=400,
                                      chunk_overlap=200,)
esops_documents = text_splitter.transform_documents(docs)
print(f"number of chunks in barbie documents : {len(esops_documents)}")

store = LocalFileStore("./cache/")
model_name = "sentence-transformers/all-mpnet-base-v2"
model_kwargs = {'device': 'cuda'}
encode_kwargs = {'normalize_embeddings': False}
core_embeddings_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs)
embedder = CacheBackedEmbeddings.from_bytes_store(core_embeddings_model,
                                                  store,
                                                  namespace=model_name)
# Create VectorStore
vectorstore = FAISS.from_documents(esops_documents,embedder)

bm25_retriever = BM25Retriever.from_documents(esops_documents)
bm25_retriever.k=2

faiss_retriever = vectorstore.as_retriever(search_kwargs={"k":2})
ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever,faiss_retriever],
                                       weights=[0.5,0.5])

from langchain.callbacks.base import BaseCallbackHandler
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from queue import Queue, Empty
from langchain.llms import  LlamaCpp

class QueueCallback(BaseCallbackHandler):
    """Callback handler for streaming LLM responses to a queue."""

    def __init__(self, q):
        self.q = q

    def on_llm_new_token(self, token: str, **kwargs: any) -> None:
        self.q.put(token)

    def on_llm_end(self, *args, **kwargs: any) -> None:
        return self.q.empty()

!CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install --upgrade --force-reinstall llama-cpp-python==0.1.83 --no-cache-dir --verbose

!gdown --id 1hVEpXeovQU-OYD6HKT48Ux8R5u1kdXB5

q = Queue()
llm = LlamaCpp(model_path="/content/llama-2-7b.Q4_0.gguf", n_ctx=2048, max_tokens=2048,  callbacks=[QueueCallback(q)], n_gpu_layers=20, n_batch=512, verbose=True)
template = """Use the following pieces of context to answer the question at the end. If you don't know the answer,\
just say that you don't know, don't try to make up an answer.

Context: {context}
Question: {question}
Helpful Answer:"""
prompt = PromptTemplate(input_variables=["context", "question"], template=template)
memory = ConversationBufferMemory(input_key="question", memory_key="history")



qa = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=ensemble_retriever,
    chain_type_kwargs={"prompt": prompt, "memory": memory},
    return_source_documents=True,
)

query = input("\nEnter a query: ")
# Get the answer from the chain
res = qa(query)
answer = res["result"]

# Print the result
print("\n\n> Question:")
print(query)
print("\n> Answer:")
print(answer)

!pip show llama-cpp-python

llm = LlamaCpp(model_path="/content/llama-2-7b.Q4_0.gguf", n_ctx=2048, max_tokens=2048, n_batch=512, verbose=False)

from huggingface_hub import hf_hub_download
model_name_or_path = "TheBloke/Llama-2-13B-Chat-GGUF"
model_basename = "llama-2-13b.ggmlv3.q4_0.bin"
model_path = hf_hub_download(repo_id=model_name_or_path, filename=model_basename)

n_gpu_layers = 40
n_batch = 256

# Loading model,
llm = LlamaCpp(
    model_path=model_path,
    max_tokens=256,
    n_gpu_layers=n_gpu_layers,
    n_batch=n_batch,
    n_ctx=1024,
    verbose=True,
)

import re
from fastapi import FastAPI
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import ElasticVectorSearch
from langchain.embeddings import HuggingFaceEmbeddings, CacheBackedEmbeddings
from langchain.llms import  LlamaCpp
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
import tiktoken
from queue import Queue, Empty
from langchain.chains import ConversationChain
from threading import Thread
from collections.abc import Generator
import gradio as gr
from langchain.document_loaders import PyMuPDFLoader,DirectoryLoader
from langchain.storage import LocalFileStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.retrievers import BM25Retriever,EnsembleRetriever


from langchain.callbacks.base import BaseCallbackHandler

class QueueCallback(BaseCallbackHandler):
    """Callback handler for streaming LLM responses to a queue."""

    def __init__(self, q):
        self.q = q

    def on_llm_new_token(self, token: str, **kwargs: any) -> None:
        self.q.put(token)

    def on_llm_end(self, *args, **kwargs: any) -> None:
        return self.q.empty()

def stream(input_text) -> Generator:
    # Create a Queue
    q = Queue()
    job_done = object()
    dir_loader = DirectoryLoader(r"/content/data",
                                 glob="*.pdf",
                                 loader_cls=PyMuPDFLoader)
    docs = dir_loader.load()
    text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=450,
                                      chunk_overlap=200,)
    esops_documents = text_splitter.transform_documents(docs)
    print(f"number of chunks in documents : {len(esops_documents)}")

    store = LocalFileStore("./cache/")
    model_name = "sentence-transformers/all-mpnet-base-v2"
    model_kwargs = {'device': 'cuda'}
    encode_kwargs = {'normalize_embeddings': False}
    core_embeddings_model = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs)
    embedder = CacheBackedEmbeddings.from_bytes_store(core_embeddings_model,
                                                      store,
                                                      namespace=model_name)
    # Create VectorStore
    vectorstore = FAISS.from_documents(esops_documents,embedder)
    bm25_retriever = BM25Retriever.from_documents(esops_documents)
    bm25_retriever.k=2

    faiss_retriever = vectorstore.as_retriever(search_kwargs={"k":2})
    ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever,faiss_retriever],
                                           weights=[0.5,0.5])

    """Logic for loading the chain you want to use should go here."""
    llm = LlamaCpp(model_path=r"/content/llama-2-7b.Q4_0.gguf", n_ctx=2048, max_tokens=2048, callbacks=[QueueCallback(q)], n_batch=512, n_gpu_layers=20, verbose=True)
    template = """Use the following pieces of context to answer the question at the end. If you don't know the answer,\
    just say that you don't know, don't try to make up an answer.

    {context}

    {history}
    Question: {question}
    Helpful Answer:"""
    prompt = PromptTemplate(input_variables=["history", "context", "question"], template=template)
    memory = ConversationBufferMemory(input_key="question", memory_key="history")



    qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=ensemble_retriever,
        chain_type_kwargs={"prompt": prompt, "memory": memory},

    )
    # Create a funciton to call - this will run in a thread
    def task():
        resp = qa.run(input_text)
        q.put(job_done)

    # Create a thread and start the function
    t = Thread(target=task)
    t.start()

    content = ""

    # Get each new token from the queue and yield for our generator
    while True:
        try:
            next_token = q.get(True, timeout=1)
            if next_token is job_done:
                break
            content += next_token
            yield next_token, content
        except Empty:
            continue
def main():

    def ask_llm(message, history):
        for next_token, content in stream(message):
            print(next_token)
            yield(content)

    demo = chatInterface = gr.ChatInterface(
                    fn=ask_llm,
    )
    demo.queue().launch()

if __name__ == "__main__":
    main()