run_localGPT_API.py

from flask import Flask, jsonify, request
from werkzeug.utils import secure_filename
import ssl
import os
import subprocess
import shutil

import torch
import logging
from langchain.chains import RetrievalQA
from langchain.embeddings import HuggingFaceInstructEmbeddings
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.llms import HuggingFacePipeline
# from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.vectorstores import Chroma
from transformers import LlamaForCausalLM, LlamaTokenizer, pipeline
from transformers import AutoTokenizer, AutoModelForCausalLM
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from constants import CHROMA_SETTINGS, PERSIST_DIRECTORY, ROOT_DIRECTORY
from transformers import GenerationConfig

DEVICE_TYPE = 'cuda'
SHOW_SOURCES = True
logging.info(f'Running on: {DEVICE_TYPE}')
logging.info(f'Display Source Documents set to: {SHOW_SOURCES}')

EMBEDDINGS = HuggingFaceInstructEmbeddings(
    model_name="hkunlp/instructor-large", model_kwargs={"device": DEVICE_TYPE}
)

# uncomment the following line if you used HuggingFaceEmbeddings in the ingest.py
# EMBEDDINGS = HuggingFaceEmbeddings(model_name='all-MiniLM-L6-v2')
if os.path.exists(PERSIST_DIRECTORY):
    try:
        shutil.rmtree(PERSIST_DIRECTORY)
    except OSError as e:
        print(f"Error: {e.filename} - {e.strerror}.")
else:
    print("The directory does not exist")

result = subprocess.run(['python', 'ingest.py'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
if result.returncode != 0:
    raise FileNotFoundError('No files were found inside SOURCE_DOCUMENTS, please put a starter file inside before starting the API!')

# load the vectorstore
DB = Chroma(
    persist_directory=PERSIST_DIRECTORY,
    embedding_function=EMBEDDINGS,
    client_settings=CHROMA_SETTINGS,
)

RETRIEVER = DB.as_retriever()

# load the LLM for generating Natural Language responses
def load_model(device_type, model_id, model_basename=None):
    """
    Select a model for text generation using the HuggingFace library.
    If you are running this for the first time, it will download a model for you.
    subsequent runs will use the model from the disk.

    Args:
        device_type (str): Type of device to use, e.g., "cuda" for GPU or "cpu" for CPU.
        model_id (str): Identifier of the model to load from HuggingFace's model hub.
        model_basename (str, optional): Basename of the model if using quantized models. 
            Defaults to None.

    Returns:
        HuggingFacePipeline: A pipeline object for text generation using the loaded model.

    Raises:
        ValueError: If an unsupported model or device type is provided.
    """
    
    logging.info(f'Loading Model: {model_id}, on: {device_type}')
    logging.info('This action can take a few minutes!')

    if model_basename is not None:
        # The code supports all huggingface models that ends with GPTQ and have some variation of .no-act.order or .safetensors in their HF repo.
        print('Using AutoGPTQForCausalLM for quantized models')

        if '.safetensors' in model_basename:
            # Remove the ".safetensors" ending if present
            model_basename = model_basename.replace('.safetensors', "")

        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
        logging.info('Tokenizer loaded')

        model = AutoGPTQForCausalLM.from_quantized(
            model_id,
            model_basename=model_basename,
            use_safetensors=True,
            trust_remote_code=True,
            device="cuda:0",
            use_triton=False,
            quantize_config=None
        )
    elif device_type.lower() == 'cuda': # The code supports all huggingface models that ends with -HF or which have a .bin file in their HF repo.
        print('Using AutoModelForCausalLM for full models')
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        logging.info('Tokenizer loaded')

        model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map='auto',
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
        model.tie_weights()
    else:
        print('Using LlamaTokenizer')
        tokenizer = LlamaTokenizer.from_pretrained(model_id)
        model = LlamaForCausalLM.from_pretrained(model_id)

    # Load configuration from the model to avoid warnings
    generation_config = GenerationConfig.from_pretrained(model_id)
    # see here for details: https://huggingface.co/docs/transformers/main_classes/text_generation#transformers.GenerationConfig.from_pretrained.returns

    # Create a pipeline for text generation
    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_length=2048,
        temperature=0,
        top_p=0.95,
        repetition_penalty=1.15,
        generation_config=generation_config
    )

    local_llm = HuggingFacePipeline(pipeline=pipe)
    logging.info('Local LLM Loaded')

    return local_llm

# for HF models
# model_id = "TheBloke/vicuna-7B-1.1-HF"
# model_id = "TheBloke/Wizard-Vicuna-7B-Uncensored-HF"
# model_id = "TheBloke/guanaco-7B-HF"
# model_id = 'NousResearch/Nous-Hermes-13b' # Requires ~ 23GB VRAM. Using STransformers alongside will 100% create OOM on 24GB cards. 
# LLM = load_model(device_type=DEVICE_TYPE, model_id=model_id)

# for GPTQ (quantized) models
# model_id = "TheBloke/Nous-Hermes-13B-GPTQ"
# model_basename = "nous-hermes-13b-GPTQ-4bit-128g.no-act.order"
# model_id = "TheBloke/WizardLM-30B-Uncensored-GPTQ"
# model_basename = "WizardLM-30B-Uncensored-GPTQ-4bit.act-order.safetensors" # Requires ~21GB VRAM. Using STransformers alongside can potentially create OOM on 24GB cards.
# model_id = "TheBloke/wizardLM-7B-GPTQ"
# model_basename = "wizardLM-7B-GPTQ-4bit.compat.no-act-order.safetensors"
model_id = "TheBloke/WizardLM-7B-uncensored-GPTQ"
model_basename = "WizardLM-7B-uncensored-GPTQ-4bit-128g.compat.no-act-order.safetensors"
LLM = load_model(device_type=DEVICE_TYPE, model_id=model_id, model_basename = model_basename)

QA = RetrievalQA.from_chain_type(
    llm=LLM, chain_type="stuff", retriever=RETRIEVER, return_source_documents=SHOW_SOURCES
)

app = Flask(__name__)
@app.route('/api/delete_source', methods=['GET'])
def delete_source_route():
    folder_name = 'SOURCE_DOCUMENTS'
    
    if os.path.exists(folder_name):
        shutil.rmtree(folder_name)
    
    os.makedirs(folder_name)

    return jsonify({"message": f"Folder '{folder_name}' successfully deleted and recreated."})



@app.route('/api/save_document', methods=['GET', 'POST'])
def save_document_route():
    if 'document' not in request.files:
        return 'No document part', 400
    file = request.files['document']
    if file.filename == '':
        return 'No selected file', 400
    if file:
        filename = secure_filename(file.filename)
        folder_path = 'SOURCE_DOCUMENTS'
        if not os.path.exists(folder_path):
            os.makedirs(folder_path)
        file_path = os.path.join(folder_path, filename)
        file.save(file_path)
        return 'File saved successfully', 200

@app.route('/api/run_ingest', methods=['GET'])
def run_ingest_route():
    global DB
    global RETRIEVER
    global QA
    try:
        if os.path.exists(PERSIST_DIRECTORY):
            try:
                shutil.rmtree(PERSIST_DIRECTORY)
            except OSError as e:
                print(f"Error: {e.filename} - {e.strerror}.")
        else:
            print("The directory does not exist")

        result = subprocess.run(['python', 'ingest.py'], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode != 0:
            return 'Script execution failed: {}'.format(result.stderr.decode('utf-8')), 500
        # load the vectorstore
        DB = Chroma(
            persist_directory=PERSIST_DIRECTORY,
            embedding_function=EMBEDDINGS,
            client_settings=CHROMA_SETTINGS,
        )
        RETRIEVER = DB.as_retriever()

        QA = RetrievalQA.from_chain_type(
            llm=LLM, chain_type="stuff", retriever=RETRIEVER, return_source_documents=SHOW_SOURCES
        )
        return 'Script executed successfully: {}'.format(result.stdout.decode('utf-8')), 200
    except Exception as e:
        return 'Error occurred: {}'.format(str(e)), 500

@app.route('/api/prompt_route', methods=['GET', 'POST'])
def prompt_route():
    global QA
    user_prompt = request.form.get('user_prompt')
    if user_prompt:
        # print(f'User Prompt: {user_prompt}')
        # Get the answer from the chain
        res = QA(user_prompt)
        answer, docs = res["result"], res["source_documents"]

        prompt_response_dict = {'Prompt': user_prompt,
                                'Answer': answer,
                                }
        
        prompt_response_dict['Sources'] = []
        for document in docs:
            prompt_response_dict['Sources'].append((os.path.basename(str(document.metadata["source"])), str(document.page_content)))

        return jsonify(prompt_response_dict), 200
    else:
        return 'No user prompt received', 400

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s',
                    level=logging.INFO)
    app.run(debug=False, port=5110)