app.py

"""
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") 
huffingface_read_token = os.getenv("huffingface_read_token")

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

#collection_name, vector_db, btn, llm_btn, slider_temperature, slider_maxtokens, slider_topk]
def build_qa_chain(collection_name, vector_db, file: str, llm_option, temperature, max_tokens, top_k, progress=gr.Progress()):

    if file == None or not file:
        raise gr.Error("Please upload a PDF first!")
        
    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:]
    collection_name = collection_name.replace("/","_" )
    
    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_model = list_llm[llm_option]
    task="text-generation"  # Explicitly specify task

    if llm_model == "mistralai/Mixtral-8x7B-Instruct-v0.1":
        llm = HuggingFaceEndpoint(
            repo_id=llm_model, 
            huggingfacehub_api_token = huffingface_read_token,
            task=task,  # Explicitly specify task
            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "load_in_8bit": True}
            temperature = temperature,
            max_new_tokens = max_tokens,
            top_k = top_k,
            load_in_8bit = True,
        )
    elif llm_model in ["HuggingFaceH4/zephyr-7b-gemma-v0.1","mosaicml/mpt-7b-instruct"]:
        raise gr.Error("LLM model is too large to be loaded automatically on free inference endpoint")
        llm = HuggingFaceEndpoint(
            repo_id=llm_model, 
            task=task,  # Explicitly specify task
            temperature = temperature,
            max_new_tokens = max_tokens,
            top_k = top_k,
        )
    elif llm_model == "microsoft/phi-2":
        # raise gr.Error("phi-2 model requires 'trust_remote_code=True', currently not supported by langchain HuggingFaceHub...")
        llm = HuggingFaceEndpoint(
            repo_id=llm_model, 
            task=task,  # Explicitly specify task
            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "trust_remote_code": True, "torch_dtype": "auto"}
            temperature = temperature,
            max_new_tokens = max_tokens,
            top_k = top_k,
            trust_remote_code = True,
            torch_dtype = "auto",
        )
    elif llm_model == "TinyLlama/TinyLlama-1.1B-Chat-v1.0":
        llm = HuggingFaceEndpoint(
            repo_id=llm_model, 
            task=task,  # Explicitly specify task
            # model_kwargs={"temperature": temperature, "max_new_tokens": 250, "top_k": top_k}
            temperature = temperature,
            max_new_tokens = 250,
            top_k = top_k,
        )
    elif llm_model == "meta-llama/Llama-2-7b-chat-hf":
        raise gr.Error("Llama-2-7b-chat-hf model requires a Pro subscription...")
        llm = HuggingFaceEndpoint(
            repo_id=llm_model,
            task=task,  # Explicitly specify task
            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k}
            temperature = temperature,
            max_new_tokens = max_tokens,
            top_k = top_k,
        )
    else:
        llm = HuggingFaceEndpoint(
            repo_id=llm_model, 
            task=task,  # Explicitly specify task
            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k, "trust_remote_code": True, "torch_dtype": "auto"}
            # model_kwargs={"temperature": temperature, "max_new_tokens": max_tokens, "top_k": top_k}
            temperature = temperature,
            max_new_tokens = max_tokens,
            top_k = top_k,
        )


    """
    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)    
        raise gr.Error("Please initialize the Chain first!")
    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,  show_error=True)
#demo.launch(share=True)