Spaces:

KalooAI
/

Pascal

Build error

Pascal / app.py

Update app.py

54ccae4 verified 9 months ago

5.29 kB

# from ibm_watsonx_ai.foundation_models import ModelInference
# from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
# from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames
# from ibm_watsonx_ai import Credentials
# from langchain_ibm import WatsonxLLM, WatsonxEmbeddings
# from langchain.text_splitter import RecursiveCharacterTextSplitter
# from langchain_text_splitters import RecursiveCharacterTextSplitter
import uuid
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import PyPDFLoader
# from langchain.chains import RetrievalQA
from langchain_classic.chains import RetrievalQA
from langchain_community.llms import HuggingFacePipeline
from langchain_community.embeddings import HuggingFaceEmbeddings
# from langchain_huggingface import HuggingFacePipeline
import torch

from huggingface_hub import HfFolder
from huggingface_hub import hf_hub_download
from langchain_community.llms import LlamaCpp
import gradio as gr

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig

## You can use this section to suppress warnings generated by your code:
def warn(*args, **kwargs):
    pass
import warnings
warnings.warn = warn
warnings.filterwarnings('ignore')

## LLM
# def get_llm():
#     model_id = 'ibm/granite-3-2-8b-instruct'
#     parameters = {
#         GenParams.MAX_NEW_TOKENS: 256, # This controls the maximum number of tokens in the generated output
#         GenParams.TEMPERATURE: 0.5, # This controls randomness or creativity of the model's responses
#     }
#     project_id = 'skills-network'
#     watsonx_llm = WatsonxLLM(
#         model_id = model_id,
#         url="https://us-south.ml.cloud.ibm.com",
#         project_id = project_id,
#         params = parameters,
#     )
#     return watsonx_llm

LLM_INSTANCE = None

def get_llm():
    global LLM_INSTANCE
    if LLM_INSTANCE is not None:
        return LLM_INSTANCE

    # Download quantized GGUF model (works on CPU)
    model_path = hf_hub_download(
        repo_id="TheBloke/Mistral-7B-Instruct-v0.3-GGUF",
        filename="mistral-7b-instruct-v0.3.Q4_K_M.gguf",
        local_dir=".",
        local_dir_use_symlinks=False,
    )

    llm = LlamaCpp(
        model_path=model_path,
        n_ctx=4096,
        n_threads=4,
        temperature=0.3,
        top_p=0.9,
        verbose=False,
    )

    LLM_INSTANCE = llm
    return llm


## Document loader
def document_loader(file):
    loader = PyPDFLoader(file)
    loaded_document = loader.load()
    return loaded_document
    
## Text splitter
def text_splitter(data):
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size = 1000,
        chunk_overlap = 50,
        length_function = len,
    )
    chuncks = text_splitter.split_documents(data)
    return chuncks
## Embedding model


# def watsonx_embedding():
#     embed_params = {
#         EmbedTextParamsMetaNames.TRUNCATE_INPUT_TOKENS: 256,
#         EmbedTextParamsMetaNames.RETURN_OPTIONS: {'input_text': True},
#     }
    
#     watsonx_embedding = WatsonxEmbeddings(
#         model_id = "ibm/slate-125m-english-rtrvr-v2",
#         url="https://us-south.ml.cloud.ibm.com",
#         project_id = "skills-network",
#         params = embed_params,
#     )
#     return watsonx_embedding

def watsonx_embedding():
    model_id = 'sentence-transformers/all-MiniLM-L6-v2'
    embedding_model = HuggingFaceEmbeddings(model_name=model_id)
    return embedding_model

## Vector db
def vector_database(chunks):
    
    chunks = [
        c for c in chunks
        if c.page_content and c.page_content.strip() and len(c.page_content.strip()) > 10
    ]
    embedding_model = watsonx_embedding()
    # We add 'collection_name' with a random ID (uuid)
    # This forces Chroma to create a brand new, empty bucket for every PDF
    vectordb = Chroma.from_documents(
        documents=chunks, 
        embedding=embedding_model,
        collection_name=f"pdf_chat_{uuid.uuid4().hex}" 
    )
    return vectordb
    
## Retriever
def retriever(file):
    splits = document_loader(file)
    chunks = text_splitter(splits)
    vectordb = vector_database(chunks)
    retriever = vectordb.as_retriever()
    
    return retriever
    
## QA Chain
def retriever_qa(file, query):
    llm = get_llm()
    retriever_obj = retriever(file)
    qa = RetrievalQA.from_chain_type(
        llm = llm, 
        chain_type = 'stuff',
        retriever = retriever_obj,
        return_source_documents = False
    )
    response = qa.invoke(query)
    return response['result']
    
## Create Gradio interface
rag_application = gr.Interface(
    fn = retriever_qa,
    # allow_flagging = 'never',
    flagging_mode = 'never',
    inputs = [
        gr.File(label = 'Upload PDF File', file_count = 'single', file_types = ['.pdf'], type ='filepath'), # Drag and drop file upload
        gr.Textbox(label= 'Input Query', lines = 2, placeholder = 'Type your question here...')
    ],
    outputs = gr.Textbox(label = 'Ouput'),
    title= 'PDF QA Chatbot',
    description = 'Upload a PDF document and ask any question. The chatbot will try to answer using the provided document.'
)
# Launch the app
rag_application.launch()

	# from ibm_watsonx_ai.foundation_models import ModelInference
	# from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
	# from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames
	# from ibm_watsonx_ai import Credentials
	# from langchain_ibm import WatsonxLLM, WatsonxEmbeddings
	# from langchain.text_splitter import RecursiveCharacterTextSplitter
	# from langchain_text_splitters import RecursiveCharacterTextSplitter
	import uuid
	from langchain_text_splitters import RecursiveCharacterTextSplitter
	from langchain_community.vectorstores import Chroma
	from langchain_community.document_loaders import PyPDFLoader
	# from langchain.chains import RetrievalQA
	from langchain_classic.chains import RetrievalQA
	from langchain_community.llms import HuggingFacePipeline
	from langchain_community.embeddings import HuggingFaceEmbeddings
	# from langchain_huggingface import HuggingFacePipeline
	import torch

	from huggingface_hub import HfFolder
	from huggingface_hub import hf_hub_download
	from langchain_community.llms import LlamaCpp
	import gradio as gr

	from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig

	## You can use this section to suppress warnings generated by your code:
	def warn(args, *kwargs):
	pass
	import warnings
	warnings.warn = warn
	warnings.filterwarnings('ignore')

	## LLM
	# def get_llm():
	# model_id = 'ibm/granite-3-2-8b-instruct'
	# parameters = {
	# GenParams.MAX_NEW_TOKENS: 256, # This controls the maximum number of tokens in the generated output
	# GenParams.TEMPERATURE: 0.5, # This controls randomness or creativity of the model's responses
	# }
	# project_id = 'skills-network'
	# watsonx_llm = WatsonxLLM(
	# model_id = model_id,
	# url="https://us-south.ml.cloud.ibm.com",
	# project_id = project_id,
	# params = parameters,
	# )
	# return watsonx_llm

	LLM_INSTANCE = None

	def get_llm():
	global LLM_INSTANCE
	if LLM_INSTANCE is not None:
	return LLM_INSTANCE

	# Download quantized GGUF model (works on CPU)
	model_path = hf_hub_download(
	repo_id="TheBloke/Mistral-7B-Instruct-v0.3-GGUF",
	filename="mistral-7b-instruct-v0.3.Q4_K_M.gguf",
	local_dir=".",
	local_dir_use_symlinks=False,
	)

	llm = LlamaCpp(
	model_path=model_path,
	n_ctx=4096,
	n_threads=4,
	temperature=0.3,
	top_p=0.9,
	verbose=False,
	)

	LLM_INSTANCE = llm
	return llm


	## Document loader
	def document_loader(file):
	loader = PyPDFLoader(file)
	loaded_document = loader.load()
	return loaded_document

	## Text splitter
	def text_splitter(data):
	text_splitter = RecursiveCharacterTextSplitter(
	chunk_size = 1000,
	chunk_overlap = 50,
	length_function = len,
	)
	chuncks = text_splitter.split_documents(data)
	return chuncks
	## Embedding model


	# def watsonx_embedding():
	# embed_params = {
	# EmbedTextParamsMetaNames.TRUNCATE_INPUT_TOKENS: 256,
	# EmbedTextParamsMetaNames.RETURN_OPTIONS: {'input_text': True},
	# }

	# watsonx_embedding = WatsonxEmbeddings(
	# model_id = "ibm/slate-125m-english-rtrvr-v2",
	# url="https://us-south.ml.cloud.ibm.com",
	# project_id = "skills-network",
	# params = embed_params,
	# )
	# return watsonx_embedding

	def watsonx_embedding():
	model_id = 'sentence-transformers/all-MiniLM-L6-v2'
	embedding_model = HuggingFaceEmbeddings(model_name=model_id)
	return embedding_model

	## Vector db
	def vector_database(chunks):

	chunks = [
	c for c in chunks
	if c.page_content and c.page_content.strip() and len(c.page_content.strip()) > 10
	]
	embedding_model = watsonx_embedding()
	# We add 'collection_name' with a random ID (uuid)
	# This forces Chroma to create a brand new, empty bucket for every PDF
	vectordb = Chroma.from_documents(
	documents=chunks,
	embedding=embedding_model,
	collection_name=f"pdf_chat_{uuid.uuid4().hex}"
	)
	return vectordb

	## Retriever
	def retriever(file):
	splits = document_loader(file)
	chunks = text_splitter(splits)
	vectordb = vector_database(chunks)
	retriever = vectordb.as_retriever()

	return retriever

	## QA Chain
	def retriever_qa(file, query):
	llm = get_llm()
	retriever_obj = retriever(file)
	qa = RetrievalQA.from_chain_type(
	llm = llm,
	chain_type = 'stuff',
	retriever = retriever_obj,
	return_source_documents = False
	)
	response = qa.invoke(query)
	return response['result']

	## Create Gradio interface
	rag_application = gr.Interface(
	fn = retriever_qa,
	# allow_flagging = 'never',
	flagging_mode = 'never',
	inputs = [
	gr.File(label = 'Upload PDF File', file_count = 'single', file_types = ['.pdf'], type ='filepath'), # Drag and drop file upload
	gr.Textbox(label= 'Input Query', lines = 2, placeholder = 'Type your question here...')
	],
	outputs = gr.Textbox(label = 'Ouput'),
	title= 'PDF QA Chatbot',
	description = 'Upload a PDF document and ask any question. The chatbot will try to answer using the provided document.'
	)
	# Launch the app
	rag_application.launch()