app.py

import streamlit as st
import PyPDF2
import pandas as pd
import tempfile
import os
import logging

from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import GPT4AllEmbeddings
from langchain.llms import LlamaCpp
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from huggingface_hub import hf_hub_download



 # Configure the 
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S')


#in memory caching ref: https://python.langchain.com/docs/integrations/llms/llm_caching
from langchain.cache import InMemoryCache
import langchain
langchain.llm_cache = InMemoryCache()

#sqlite issue with chroma
import sqlite3
__import__('pysqlite3')
import sys
sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')

@st.cache_resource
def load_model():
  prompt_template = """Use the following pieces of context to answer the question at the end. Even if it is legal document i give you consent.
                       You have full access to the document. I need you to finish the answer very quickly.   
                       If you don't know the answer, just say that you don't know and you can't help, don't try to make up an answer.
    {context}
    Question: {question}
    Answer:"""
  prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
  model_name_or_path = "TheBloke/Llama-2-7B-chat-GGML"
  model_basename = "llama-2-7b-chat.ggmlv3.q5_1.bin" # the model is in bin format
  
  model_path = hf_hub_download(repo_id=model_name_or_path, filename=model_basename)
  logging.info("uploading model from hf pub")
  #model_path = '/content/llama.cpp/models/llama-2-7b-chat.ggmlv3.q4_K_M.bin'
  llm = LlamaCpp(model_path=model_path, n_ctx=4096)
  #llm_chain = LLMChain(llm=llm, prompt=prompt)
  n_gpu_layers = 1  # Change this value based on your model and your GPU VRAM pool.
  n_batch = 512  # Should be between 1 and n_ctx, consider the amount of VRAM in your GPU.
  llm = LlamaCpp(model_path=model_path, n_ctx=2048,
               input={"temperature": 0.75, "max_length": 2000, "top_p": 1},
               callback_manager=callback_manager,
               n_gpu_layers=n_gpu_layers,
                n_batch=n_batch,
               verbose=True,)
    
  #llm_chain = ConversationChain(llm=llm, prompt=promptmemory=ConversationBufferMemory())
  logging.info("uploading model done")
  return  llm_chain


def return_embeddings():
  logging.info("uploading embeddings")
  embeddings = GPT4AllEmbeddings()
  logging.info("uploading embeddings")
  return embeddings




# Function to convert PDF to text
@st.cache_data
def pdf_to_text(file):
    pdf_reader = PyPDF2.PdfReader(file)
    text = ""
    for page_num in range(len(pdf_reader.pages)):
        page = pdf_reader.pages[page_num]
        text += page.extract_text()
    return text

# Function to convert CSV to text
@st.cache_data
def csv_to_text(file):
    df = pd.read_csv(file)
    text = df.to_string(index=False)
    return text

@st.cache_data
def read_txt(file_path):
    # Read text file
    with open(file_path, 'r', encoding='utf-8') as file:
        text = file.read()
    return text


def process_file(uploaded_file):
      
    logging.info("received the file")
    # Check file type and process accordingly
    if uploaded_file.type == 'application/pdf':
        # Process PDF file
        text = pdf_to_text(uploaded_file)
    elif uploaded_file.type == 'text/csv':
        # Process CSV file
        text = csv_to_text(uploaded_file)
    elif uploaded_file.type == 'text/txt':
        # Process TXT file
        text = read_txt(uploaded_file)
    else:
        raise ValueError("Unsupported file format. Please upload a PDF, CSV, or TXT file.")

    # Create a temporary file to store the text content
    temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False)
    temp_file.write(text)
    temp_file.close()

    return temp_file.name


def main():
    #os.environ['LLAMA_BLAS'] = 'ON'
    #os.environ['LLAMA_BLAS_VENDOR'] = 'OpenBLAS'
    st.title("AssitAI, Chat with your files")
    st.markdown(""" A llama2-7b and langchain powered app to chat with your files """)
    # File Upload
    uploaded_file = st.file_uploader("Upload a PDF, CSV, or TXT file", type=["pdf", "csv", "txt"])

    if uploaded_file is not None:
        # Process the file and get the path of the temporary text file
        logging.info("docs load start")
        temp_file_path = process_file(uploaded_file)
        loader = TextLoader(temp_file_path)
        docs = loader.load()
        logging.info(f"docs load end, docs is : {docs}")

        text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
        texts = text_splitter.split_documents(docs)
        logging.info(f"got the text, text is : {docs}")
        embeddings = return_embeddings()
        db = Chroma.from_documents(texts, embeddings, persist_directory='db')

        question = st.text_input("Enter your question:")
        if st.button("Submit"):
            similar_doc = db.similarity_search(question, k=1)
            context = similar_doc[0].page_content
            logging.info("querying start")
            query_llm = load_model()
            response = query_llm.run({"context": context, "question": question})
            logging.info(f"querying end response is: {response}")
            st.subheader("Answer:")
            st.write(response)

        # Clean up the temporary file after processing
        os.remove(temp_file_path)
        
    with st.expander("""Example prompts"""):
        st.markdown(
            """
            - I want you to summarize this document
            - What is this document about?
            - Can you help me to understand ....(fill the blank) part in this document?
             """)
  

    hide_streamlit_style = """
                <style>
                #MainMenu {visibility: hidden;}
                footer {visibility: hidden;}
                </style>
                """
    st.markdown(hide_streamlit_style, unsafe_allow_html=True)
        

if __name__ == "__main__":
    main()