Create app.py

app.py

@@ -0,0 +1,472 @@
+"""
+This code uses the PyMuPDF package.
+
+PyMuPDF is AGPL licensed, please refer to:
+https://pymupdf.readthedocs.io/en/latest/about.html#license-and-copyright
+"""
+
+"""
+Code below is based on an implementation by Sunil Kumar Dash:
+
+MIT License
+
+Copyright (c) 2023 Sunil Kumar Dash
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+from huggingface_hub import InferenceClient
+from langchain_openai import AzureOpenAIEmbeddings
+#from langchain_community.chat_models import AzureChatOpenAI
+from langchain_openai import AzureChatOpenAI
+
+
+
+
+
+from typing import Any
+import gradio as gr
+from langchain_openai import OpenAIEmbeddings
+from langchain_community.vectorstores import Chroma
+import chromadb
+#to handle the tenant issue
+chromadb.api.client.SharedSystemClient.clear_system_cache()
+
+from langchain.chains import ConversationalRetrievalChain
+from langchain_openai import ChatOpenAI
+
+from langchain_community.document_loaders import PyMuPDFLoader
+from langchain.schema.document import Document
+
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.memory import ConversationBufferMemory
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain_community.embeddings import HuggingFaceEmbeddings 
+
+
+# for hugging face llm
+from transformers import AutoTokenizer
+import transformers
+import torch
+import tqdm 
+import accelerate
+
+
+import pymupdf
+from PIL import Image
+import os
+import re
+import uuid
+
+import os
+import wget
+import subprocess
+import urllib.request
+import requests
+
+from pathlib import Path
+from unidecode import unidecode
+
+api_key = os.getenv("OPENAI_API_KEY")
+user_agent = os.getenv("USER_AGENT")
+
+dr_link_url1 = os.getenv("DR_LINK_1")
+dr_link_url2 = os.getenv("DR_LINK_2") 
+azure_endpt = os.getenv("AZURE_ENDPT") 
+
+
+list_llm = ["mistralai/Mistral-7B-Instruct-v0.2", "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.1", \
+    "google/gemma-7b-it","google/gemma-2b-it", \
+    "HuggingFaceH4/zephyr-7b-beta", "HuggingFaceH4/zephyr-7b-gemma-v0.1", \
+    "meta-llama/Llama-2-7b-chat-hf", "microsoft/phi-2", \
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "mosaicml/mpt-7b-instruct", "tiiuae/falcon-7b-instruct", \
+    "google/flan-t5-xxl"
+]
+list_llm_simple = [os.path.basename(llm) for llm in list_llm]
+
+
+"""
+enable_box = gr.Textbox(
+    value=None, placeholder="Upload your OpenAI API key", interactive=True
+)
+disable_box = gr.Textbox(value="OpenAI API key is set", interactive=False)
+"""
+
+def set_apikey(api_key: str):
+    print("API Key set")
+    app.OPENAI_API_KEY = api_key
+    #return disable_box
+
+"""
+def enable_api_box():
+    return enable_box
+"""
+
+def add_text(history, text: str):
+    if not text:
+        raise gr.Error("enter text")
+    print("in add_text history="+str(history))
+    print("in add_text text="+str(text))
+    
+    history = history + [(text, "")]
+    return history
+
+
+class my_app:
+    def __init__(self, OPENAI_API_KEY: str = None) -> None:
+        print("init")
+        self.OPENAI_API_KEY: str = api_key
+        #self.chain = None
+        #self.chat_history: list = []
+        #self.N: int = 0
+        self.count: int = 0
+
+    def __call__(self, file: str) -> Any:
+        print("call")
+        #if self.count == 0:
+            #vincent added
+            #self.chain = None
+            #self.chat_history: list = []
+            #self.N: int = 0
+            #self.count: int = 0
+            
+            #self.chain = self.build_chain(file)
+            #self.count += 1
+            
+        #vincent added
+        
+        #else:
+            #self.chain = self.build_chain(file)
+            #self.count += 1
+            
+        #return self.chain
+    
+
+def process_file2(file: str):
+    
+    loader = PyMuPDFLoader(file.name)
+    documents = loader.load()
+    pattern = r"/([^/]+)$"
+    match = re.search(pattern, file.name)
+    try:
+        file_name = match.group(1)
+    except:
+        file_name = os.path.basename(file)
+        
+    return documents, file_name
+
+def create_collection_name(filepath):
+    # Extract filename without extension
+    collection_name = Path(filepath).stem
+    # Fix potential issues from naming convention
+    ## Remove space
+    collection_name = collection_name.replace(" ","-") 
+    ## ASCII transliterations of Unicode text
+    collection_name = unidecode(collection_name)
+    ## Remove special characters
+    #collection_name = re.findall("[\dA-Za-z]*", collection_name)[0]
+    collection_name = re.sub('[^A-Za-z0-9]+', '-', collection_name)
+    ## Limit length to 50 characters
+    collection_name = collection_name[:50]
+    ## Minimum length of 3 characters
+    if len(collection_name) < 3:
+        collection_name = collection_name + 'xyz'
+    ## Enforce start and end as alphanumeric character
+    if not collection_name[0].isalnum():
+        collection_name = 'A' + collection_name[1:]
+    if not collection_name[-1].isalnum():
+        collection_name = collection_name[:-1] + 'Z'
+    print('Filepath: ', filepath)
+    print('Collection name: ', collection_name)
+    return collection_name
+
+def build_qa_chain(collection_name, vector_db, file: str):
+    print("in build_qa_chain="+file.name) 
+    documents, file_name = process_file2(file)
+    # Load embeddings model
+    #embeddings = OpenAIEmbeddings(openai_api_key=self.OPENAI_API_KEY)
+
+    #vincent for old LLM
+    """
+    embeddings = AzureOpenAIEmbeddings(
+    model="text-embedding-ada-002",
+    # dimensions: Optional[int] = None, # Can specify dimensions with new text-embedding-3 models   
+    azure_endpoint=azure_endpt , # If not provided, will read env variable AZURE_OPENAI_ENDPOINT
+    openai_api_key=api_key, # Can provide an API key directly. If missing read env variable AZURE_OPENAI_API_KEY   
+    #openai_api_version="2023-05-15", # If not provided, will read env variable AZURE_OPENAI_API_VERSION
+    openai_api_version="2023-05-15", # If not provided, will read env variable AZURE_OPENAI_API_VERSION
+    )
+    """    
+    
+    #vincent for new LLM
+    embeddings = HuggingFaceEmbeddings()
+    
+    
+    #vincent added to handle the tenant problem 20250211
+    chromadb.api.client.SharedSystemClient.clear_system_cache()   
+    new_client = chromadb.EphemeralClient()
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key='answer',
+        return_messages=True
+    )
+    
+    # added by vincent
+    text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=10)
+    chunked_documents = text_splitter.split_documents(documents)
+
+    #list_file_path = [x.name for x in list_file_obj if x is not None]
+    list_file_path = file.name 
+    # Create collection_name for vector database
+    # vincent fix InvalidCollectionException 20250212
+    #collection_name = create_collection_name(list_file_path[0])
+    collection_name = "pdf_docs_l_"+file.name[-10:]
+    
+    vector_db = Chroma.from_documents(
+            documents=chunked_documents,
+            embedding=embeddings,
+            client=new_client,
+            #collection_name=file_name,
+            #persist_directory = "db_" + file_name,
+            collection_name=collection_name,
+        )
+    """
+    chain = ConversationalRetrievalChain.from_llm(
+    ChatOpenAI(temperature=0.0, openai_api_key=self.OPENAI_API_KEY),
+    retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}),
+    return_source_documents=True,
+    )
+    """
+
+    #vincent added for old LLM
+    """
+    chain = ConversationalRetrievalChain.from_llm(
+    #ChatOpenAI(temperature=0.0, openai_api_key=self.OPENAI_API_KEY),
+    
+    AzureChatOpenAI(
+                temperature=0.0, openai_api_key=api_key, api_version="2024-08-01-preview", 
+                model_name="gpt-4o", azure_endpoint=azure_endpt),
+    #vincent modified
+    retriever=vector_db.as_retriever(),
+    #retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}),            
+    return_source_documents=True,
+    chain_type="stuff",     
+    memory=memory,
+    )
+    
+    """
+    #vincent for new LLM
+    #llm_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    #llm_model = "meta-llama/Llama-2-7b-chat-hf"
+    llm = HuggingFaceEndpoint(
+            repo_id=llm_model, 
+            task="text-generation",  # Explicitly specify task
+            # model_kwargs={"temperature": temperature, "max_new_tokens": 250, "top_k": top_k}
+            temperature = 0.01,
+            max_new_tokens = 250,
+            top_k = 3,
+        )
+    
+    chain = ConversationalRetrievalChain.from_llm(
+        llm,
+        retriever=vector_db.as_retriever(),
+        chain_type="stuff", 
+        memory=memory,
+        # combine_docs_chain_kwargs={"prompt": your_prompt})
+        return_source_documents=True,
+        #return_generated_question=False,
+        verbose=False,
+    )
+    
+    #vincent added 20250211
+    app.count += 1
+    return collection_name, vector_db, chain, "Complete!"
+
+
+def get_response(collection_name, vector_db, qa_chain, history, query, file):
+    #vincent added
+    set_apikey(api_key)
+    #print("in get_response count=" + str(app.count))
+    if not file:
+        raise gr.Error(message="Upload a PDF")
+        
+    formatted_chat_history = list(history)
+    formatted_chat_history = formatted_chat_history[:len(formatted_chat_history)-1]
+    print("in get_response query="+ query)
+    #print("in get_response chat_history="+ str(app.chat_history))
+    print("in get_response formatted_chat_history="+ str(formatted_chat_history))
+    #print("in get_response history="+ str(history))
+
+    chat_history_tuples = []
+    for message in formatted_chat_history:
+        chat_history_tuples.append((message[0], message[1]))
+
+    #vincent added 20250211
+    if app.count == 0:
+        collection_name, vector_db, qa_chain = build_qa_chain(collection_name, vector_db, file)    
+    result = qa_chain.invoke(
+        {"question": query, "chat_history": chat_history_tuples}, return_only_outputs=True
+        #{"question": query, "chat_history": format_chat_history(query, history)}, return_only_outputs=True
+    )
+    
+    
+    #app.chat_history += [(query, result["answer"])]
+    ##app.N = list(result["source_documents"][0])[1][1]["page"]
+    for char in result["answer"]:
+        history[-1][-1] += char
+        yield collection_name, vector_db, qa_chain, history, ""
+
+    #print("answer:"+ result["answer"])
+
+
+def render_file(file):
+    #print("in render_file="+file.name+" count="+str(app.count))
+    doc = pymupdf.open(file.name)
+    # vincent: issue in N
+    page = doc[N]
+    # Render the page as a PNG image with a resolution of 150 DPI
+    pix = page.get_pixmap(dpi=150)
+    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    return image
+
+
+def purge_chat_and_render_first(file):
+    print("purge_chat_and_render_first")
+    # Purges the previous chat session so that the bot has no concept of previous documents
+    chat_history = []
+    history = []
+    #count = 0
+
+    #vincent added 20250211
+    #count = count + 1
+    app.count = 0
+
+    # Use PyMuPDF to render the first page of the uploaded document
+    doc = pymupdf.open(file.name)
+    page = doc[0]
+    # Render the page as a PNG image with a resolution of 150 DPI
+    pix = page.get_pixmap(dpi=150)
+    image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+    return image, []
+
+app = my_app()
+
+with gr.Blocks() as demo:
+    
+    vector_db = gr.State()
+    qa_chain = gr.State()
+    collection_name = gr.State()
+    #N = gr.Number()
+    #count = gr.Number()
+
+    N = 0
+    count = 0
+    
+    #chat_history = gr.State()
+    #chat_history: list = []
+    #chat_history = []
+    
+    with gr.Column():
+        """
+        with gr.Row():
+
+            with gr.Column(scale=1):
+                api_key = gr.Textbox(
+                    placeholder="Enter OpenAI API key and hit <RETURN>",
+                    show_label=False,
+                    interactive=True
+                )
+        """
+        with gr.Row():
+                llm_btn = gr.Radio(list_llm_simple, \
+                    label="LLM models", value = list_llm_simple[0], type="index", info="Choose your LLM model")
+
+        with gr.Accordion("Advanced options - LLM model", open=False):
+                with gr.Row():
+                    slider_temperature = gr.Slider(minimum = 0.01, maximum = 1.0, value=0.7, step=0.1, label="Temperature", info="Model temperature", interactive=True)
+                with gr.Row():
+                    slider_maxtokens = gr.Slider(minimum = 224, maximum = 4096, value=1024, step=32, label="Max Tokens", info="Model max tokens", interactive=True)
+                with gr.Row():
+                    slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="top-k samples", info="Model top-k samples", interactive=True)
+        
+        with gr.Row():
+                llm_progress = gr.Textbox(value="None",label="QA chain initialization")
+        with gr.Row():
+                qachain_btn = gr.Button("Initialize Question Answering chain")    
+            
+        with gr.Row():
+            with gr.Column(scale=2):
+                with gr.Row():
+                    chatbot = gr.Chatbot(value=[], elem_id="chatbot")
+                with gr.Row():
+                    txt = gr.Textbox(
+                        show_label=False,
+                        placeholder="Enter text and press submit",
+                        scale=2
+                    )
+                    submit_btn = gr.Button("submit", scale=1)
+
+            with gr.Column(scale=1):
+                with gr.Row():
+                    show_img = gr.Image(label="Upload PDF")
+                with gr.Row():
+                    btn = gr.UploadButton("📁 upload a PDF", file_types=[".pdf"])
+
+    """
+    api_key.submit(
+        fn=set_apikey,
+        inputs=[api_key],
+        outputs=[
+            api_key,
+        ],
+    )
+    """
+
+    btn.upload(
+        fn=purge_chat_and_render_first,
+        inputs=[btn],
+        outputs=[show_img, chatbot],
+    )
+
+    qachain_btn.click(build_qa_chain, \
+            inputs=[collection_name, vector_db, btn, llm_btn, slider_temperature, slider_maxtokens, slider_topk], \
+            outputs=[collection_name, vector_db, qa_chain, llm_progress]).then(lambda:[None], \
+            inputs=None, \
+            outputs=[chatbot], \
+            queue=False)
+
+    submit_btn.click(
+        fn=add_text,
+        inputs=[chatbot, txt],
+        outputs=[
+            chatbot,
+        ],
+        queue=False,
+    ).success(
+        fn=get_response, inputs=[collection_name,vector_db, qa_chain, chatbot, txt, btn], outputs=[collection_name,vector_db, qa_chain, chatbot, txt]
+    ).success(
+        fn=render_file, inputs=[btn], outputs=[show_img]
+    )
+
+#demo.queue()
+#demo.launch(share=True, ssr_mode=False)
+#demo.launch()
+demo.queue().launch(share=True)
+#demo.launch(share=True)