init

Dockerfile

@@ -0,0 +1,19 @@
+# Use an official Python runtime as a parent image
+FROM python:3.9-slim
+
+# Set the working directory in the container
+WORKDIR /usr/src/app
+
+
+# Install any needed packages specified in requirements.txt
+COPY requirements.txt ./
+RUN pip install -r requirements.txt
+
+# Copy the rest of your application's code
+COPY ./src .
+
+# Make port 7860 available to the world outside this container
+EXPOSE 7860
+
+# Run app.py when the container launches
+CMD ["python", "./app.py"]

requirements.txt

@@ -0,0 +1,8 @@
+docarray==0.39.1
+faiss-cpu==1.7.4
+gradio==4.8.0
+langchain==0.0.348
+openai==1.3.8
+pypdf==3.17.2
+tiktoken==0.5.2
+transformers==4.36.0

src/app.py

@@ -0,0 +1,61 @@
+import gradio as gr
+from utils import *
+
+prompt_keys = load_prompts_list_from_json('prompts.json')
+
+with gr.Blocks(gr.themes.Soft(primary_hue=gr.themes.colors.slate, secondary_hue=gr.themes.colors.purple)) as demo:
+    with gr.Row():
+
+        with gr.Column(scale=1, variant = 'panel'):
+            # gr.HTML(f"<img src='file/logo.png' width='100' height='100'>")
+            files = gr.File(type="filepath", file_count="multiple")
+            with gr.Row(equal_height=True):
+                vector_index_btn = gr.Button('Create vector store', variant='primary',scale=1)
+                vector_index_msg_out = gr.Textbox(show_label=False, lines=1,scale=1, placeholder="Creating vectore store ...")
+            
+            prompt_dropdown = gr.Dropdown(label="Select a prompt", choices=prompt_keys, value=prompt_keys[0])
+
+            with gr.Accordion(label="Text generation tuning parameters"):
+                temperature = gr.Slider(label="temperature", minimum=0.1, maximum=1, value=0.1, step=0.05)
+                max_new_tokens = gr.Slider(label="max_new_tokens", minimum=1, maximum=4096, value=1024, step=1)
+                k_context=gr.Slider(label="k_context", minimum=1, maximum=15, value=5, step=1)
+
+            vector_index_btn.click(upload_and_create_vector_store, inputs=[files], outputs=vector_index_msg_out)
+
+        with gr.Column(scale=1, variant = 'panel'):
+            with gr.Row(equal_height=True):
+
+                with gr.Column(scale=1):
+                    llm = gr.Dropdown(choices= ["gpt-3.5-turbo", "gpt-3.5-turbo-instruct", "gpt-3.5-turbo-16k", "gpt-4", "gpt-4-32k"], 
+                                       label="Select the model")
+
+                with gr.Column(scale=1):
+                    model_load_btn = gr.Button('Load model', variant='primary',scale=2)
+                    load_success_msg = gr.Textbox(show_label=False,lines=1, placeholder="Model loading ...")
+            chatbot = gr.Chatbot([], elem_id="chatbot",
+                                label='Chatbox', height=725, )
+
+            txt = gr.Textbox(label= "Question",lines=2,placeholder="Enter your question and press shift+enter ")
+
+            with gr.Row():
+
+                with gr.Column(scale=1):
+                    submit_btn = gr.Button('Submit',variant='primary', size = 'sm')
+
+                with gr.Column(scale=1):
+                    clear_btn = gr.Button('Clear',variant='stop',size = 'sm')
+
+            model_load_btn.click(load_models, [llm], load_success_msg, api_name="load_models")
+
+            txt.submit(add_text, [chatbot, txt], [chatbot, txt]).then(
+                bot, [chatbot,prompt_dropdown,temperature,max_new_tokens,k_context], chatbot)
+            submit_btn.click(add_text, [chatbot, txt], [chatbot, txt]).then(
+                bot, [chatbot,prompt_dropdown,temperature, max_new_tokens,k_context], chatbot).then(
+                    clear_cuda_cache, None, None
+                )
+
+            clear_btn.click(lambda: None, None, chatbot, queue=False)
+
+if __name__ == '__main__':
+    # demo.queue(concurrency_count=3)
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)

src/conversation.py

@@ -0,0 +1,50 @@
+from langchain.vectorstores import FAISS
+from langchain.chains import ConversationalRetrievalChain
+from langchain.chat_models import ChatOpenAI
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.prompts import PromptTemplate
+import os
+
+openai_api_key = os.environ.get("OPENAI_API_KEY")
+
+class Conversation_RAG:
+    def __init__(self, model_name="gpt-3.5-turbo"):
+        self.model_name = model_name
+
+    def create_vectordb(self):
+        vectordb = FAISS.load_local("./db/faiss_index", OpenAIEmbeddings())
+        
+        return vectordb
+    
+    def create_model(self, max_new_tokens=512, temperature=0.1):
+
+        llm = ChatOpenAI(
+            openai_api_key=openai_api_key,
+            model_name=self.model_name,
+            temperature=temperature,
+            max_tokens=max_new_tokens,
+            )
+    
+        return llm
+    
+    def create_conversation(self, model, vectordb, k_context=5, instruction="Use the following pieces of context to answer the question at the end by. Generate the answer based on the given context only. If you do not find any information related to the question in the given context, just say that you don't know, don't try to make up an answer. Keep your answer expressive."):
+        
+        print(instruction)
+
+        template = instruction + """
+        context:\n
+        {context}\n
+        data: {question}\n
+        """
+
+        QCA_PROMPT = PromptTemplate(input_variables=["instruction", "context", "question"], template=template)
+
+        qa = ConversationalRetrievalChain.from_llm(
+            llm=model,
+            chain_type='stuff',
+            retriever=vectordb.as_retriever(search_kwargs={"k": k_context}),
+            combine_docs_chain_kwargs={"prompt": QCA_PROMPT},
+            get_chat_history=lambda h: h,
+            verbose=True
+        )
+        return qa

src/setup.py

@@ -0,0 +1,16 @@
+from conversation import Conversation_RAG
+from vector_index import *
+
+class ModelSetup:
+    def __init__(self, model_name):
+
+        self.model_name = model_name
+
+    def setup(self):
+
+        conv_rag = Conversation_RAG(self.model_name)
+
+        self.vectordb = conv_rag.create_vectordb()
+        self.pipeline = conv_rag.create_model()
+
+        return "Model Setup Complete"

src/utils.py

@@ -0,0 +1,70 @@
+import gc
+from conversation import Conversation_RAG
+from vector_index import *
+from setup import ModelSetup
+import json
+
+def load_models(model_name):
+    global conv_qa 
+    conv_qa = Conversation_RAG(model_name)
+    global model_setup
+    model_setup = ModelSetup(model_name)
+    success_prompt = model_setup.setup()
+    return success_prompt
+
+def get_chat_history(inputs):
+
+    res = []
+    for human, ai in inputs:
+        res.append(f"Human:{human}\nAssistant:{ai}")
+    return "\n".join(res)
+
+def add_text(history, text):
+
+    history = history + [[text, None]]
+    return history, ""
+
+
+def bot(history,
+        instruction="Use the following pieces of context to answer the question at the end. Generate the answer based on the given context only if you find the answer in the context. If you do not find any information related to the question in the given context, just say that you don't know, don't try to make up an answer. Keep your answer expressive.",
+        temperature=0.1,
+        max_new_tokens=512,
+        k_context=5,
+        ):
+    
+    instruction = load_prompt('prompts.json', instruction)
+
+    model = conv_qa.create_model(max_new_tokens=max_new_tokens, temperature=temperature)
+                             
+    qa = conv_qa.create_conversation(
+                             model=model,
+                             vectordb=model_setup.vectordb,
+                             k_context=k_context,
+                             instruction=instruction
+    )
+
+    chat_history_formatted = get_chat_history(history[:-1])
+    res = qa(
+        {
+            'question': history[-1][0],
+            'chat_history': chat_history_formatted
+        }
+    )
+
+    history[-1][1] = res['answer']
+    return history
+
+def clear_cuda_cache():
+
+    gc.collect()
+    return None
+
+def load_prompts_list_from_json(json_filepath):
+    with open(json_filepath, 'r') as file:
+        data = json.load(file)
+    return list(data.keys())
+
+def load_prompt(json_filepath, key):
+    with open(json_filepath, 'r') as file:
+        data = json.load(file)
+    return data.get(key, key)

src/vector_index.py

@@ -0,0 +1,72 @@
+from langchain.vectorstores import FAISS
+from langchain.document_loaders.csv_loader import CSVLoader
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import OpenAIEmbeddings
+import os, shutil
+
+
+def create_vector_store_index(file_path):
+
+    file_path_split = file_path.split(".")
+    file_type = file_path_split[-1].rstrip('/')
+
+    if file_type == 'csv':
+        print(file_path)
+        loader = CSVLoader(file_path=file_path)
+        documents = loader.load()
+
+    elif file_type == 'pdf':
+        loader = PyPDFLoader(file_path)
+        pages = loader.load()
+
+        text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size = 512,
+        chunk_overlap = 128,)
+
+        documents = text_splitter.split_documents(pages)
+    
+    file_output = "./db/faiss_index"
+
+    try:
+        vectordb = FAISS.load_local(file_output, OpenAIEmbeddings())
+        vectordb.add_documents(documents)
+    except:
+        print("No vector store exists. Creating new one...")
+        vectordb = FAISS.from_documents(documents, OpenAIEmbeddings())
+
+    vectordb.save_local(file_output)
+
+    return "Vector store index is created."
+
+
+def upload_and_create_vector_store(files):
+    current_folder = os.getcwd()
+    data_folder = os.path.join(current_folder, "data")
+
+    # Create the directory if it doesn't exist
+    if not os.path.exists(data_folder):
+        os.makedirs(data_folder)
+
+    index_success_msg = "No new indices added."
+
+    for file in files:
+        # Save each file to a permanent location
+        file_path = file
+        split_file_name = file_path.split("/")
+        file_name = split_file_name[-1]
+        permanent_file_path = os.path.join(data_folder, file_name)
+
+        if os.path.exists(permanent_file_path):
+            print(f"File {file_name} already exists. Skipping.")
+            continue
+
+        shutil.copy(file, permanent_file_path)
+
+        # Access the path of the saved file
+        print(f"File saved to: {permanent_file_path}")
+
+        # Create an index for each file and store the success messages
+        index_success_msg = create_vector_store_index(permanent_file_path)
+
+    return index_success_msg