Spaces:
Runtime error
Runtime error
| import datetime | |
| import gradio as gr | |
| import time | |
| import uuid | |
| from langchain.embeddings import OpenAIEmbeddings | |
| from langchain.vectorstores import Chroma | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain.embeddings import SentenceTransformerEmbeddings | |
| import os | |
| from langchain.document_loaders import WebBaseLoader, TextLoader, Docx2txtLoader, PyMuPDFLoader, UnstructuredPowerPointLoader | |
| from whatsapp_chat_custom import WhatsAppChatLoader # use this instead of from langchain.document_loaders import WhatsAppChatLoader | |
| from collections import deque | |
| import re | |
| from bs4 import BeautifulSoup | |
| import requests | |
| from urllib.parse import urlparse | |
| import mimetypes | |
| from pathlib import Path | |
| import tiktoken | |
| import gdown | |
| from langchain.chat_models import ChatOpenAI | |
| from langchain.llms import OpenAI | |
| from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams | |
| from ibm_watson_machine_learning.foundation_models.utils.enums import DecodingMethods | |
| from ibm_watson_machine_learning.foundation_models import Model | |
| from ibm_watson_machine_learning.foundation_models.extensions.langchain import WatsonxLLM | |
| from genai.credentials import Credentials | |
| from genai.extensions.langchain import LangChainInterface | |
| from genai.schemas import GenerateParams | |
| # Regex pattern to match a URL | |
| HTTP_URL_PATTERN = r'^http[s]*://.+' | |
| mimetypes.init() | |
| media_files = tuple([x for x in mimetypes.types_map if mimetypes.types_map[x].split('/')[0] in ['image', 'video', 'audio']]) | |
| filter_strings = ['/email-protection#'] | |
| def getOaiCreds(key): | |
| key = key if key else 'Null' | |
| return {'service': 'openai', | |
| 'oai_key' : key | |
| } | |
| def getBamCreds(key): | |
| key = key if key else 'Null' | |
| return {'service': 'bam', | |
| 'bam_creds' : Credentials(key, api_endpoint=os.getenv("BAM_ENDPOINT",'https://bam-api.res.ibm.com/v1')) | |
| } | |
| def getWxCreds(key, p_id): | |
| key = key if key else 'Null' | |
| p_id = p_id if p_id else 'Null' | |
| return {'service': 'watsonx', | |
| 'credentials' : {"url": "https://us-south.ml.cloud.ibm.com", "apikey": key }, | |
| 'project_id': p_id | |
| } | |
| def getPersonalBotApiKey(): | |
| if os.getenv("OPENAI_API_KEY"): | |
| return getOaiCreds(os.getenv("OPENAI_API_KEY")) | |
| elif os.getenv("WX_API_KEY") and os.getenv("WX_PROJECT_ID"): | |
| return getWxCreds(os.getenv("WX_API_KEY"), os.getenv("WX_PROJECT_ID")) | |
| elif os.getenv("BAM_API_KEY"): | |
| return getBamCreds(os.getenv("BAM_API_KEY")) | |
| else: | |
| return {} | |
| def getOaiLlm(temp, modelNameDD, api_key_st): | |
| modelName = modelNameDD.split('(')[0].strip() | |
| # check if the input model is chat model or legacy model | |
| try: | |
| ChatOpenAI(openai_api_key=api_key_st['oai_key'], temperature=0,model_name=modelName,max_tokens=1).predict('') | |
| llm = ChatOpenAI(openai_api_key=api_key_st['oai_key'], temperature=float(temp),model_name=modelName) | |
| except: | |
| OpenAI(openai_api_key=api_key_st['oai_key'], temperature=0,model_name=modelName,max_tokens=1).predict('') | |
| llm = OpenAI(openai_api_key=api_key_st['oai_key'], temperature=float(temp),model_name=modelName) | |
| return llm | |
| MAX_NEW_TOKENS = 1024 | |
| TOP_K = None | |
| TOP_P = 1 | |
| def getWxLlm(temp, modelNameDD, api_key_st): | |
| modelName = modelNameDD.split('(')[0].strip() | |
| wxModelParams = { | |
| GenParams.DECODING_METHOD: DecodingMethods.SAMPLE, | |
| GenParams.MAX_NEW_TOKENS: MAX_NEW_TOKENS, | |
| GenParams.TEMPERATURE: float(temp), | |
| GenParams.TOP_K: TOP_K, | |
| GenParams.TOP_P: TOP_P | |
| } | |
| model = Model( | |
| model_id=modelName, | |
| params=wxModelParams, | |
| credentials=api_key_st['credentials'], project_id=api_key_st['project_id']) | |
| llm = WatsonxLLM(model=model) | |
| return llm | |
| def getBamLlm(temp, modelNameDD, api_key_st): | |
| modelName = modelNameDD.split('(')[0].strip() | |
| parameters = GenerateParams(decoding_method="sample", max_new_tokens=MAX_NEW_TOKENS, temperature=float(temp), top_k=TOP_K, top_p=TOP_P) | |
| llm = LangChainInterface(model=modelName, params=parameters, credentials=api_key_st['bam_creds']) | |
| return llm | |
| def get_hyperlinks(url): | |
| try: | |
| reqs = requests.get(url) | |
| if not reqs.headers.get('Content-Type').startswith("text/html") or 400<=reqs.status_code<600: | |
| return [] | |
| soup = BeautifulSoup(reqs.text, 'html.parser') | |
| except Exception as e: | |
| print(e) | |
| return [] | |
| hyperlinks = [] | |
| for link in soup.find_all('a', href=True): | |
| hyperlinks.append(link.get('href')) | |
| return hyperlinks | |
| # Function to get the hyperlinks from a URL that are within the same domain | |
| def get_domain_hyperlinks(local_domain, url): | |
| clean_links = [] | |
| for link in set(get_hyperlinks(url)): | |
| clean_link = None | |
| # If the link is a URL, check if it is within the same domain | |
| if re.search(HTTP_URL_PATTERN, link): | |
| # Parse the URL and check if the domain is the same | |
| url_obj = urlparse(link) | |
| if url_obj.netloc.replace('www.','') == local_domain.replace('www.',''): | |
| clean_link = link | |
| # If the link is not a URL, check if it is a relative link | |
| else: | |
| if link.startswith("/"): | |
| link = link[1:] | |
| elif link.startswith(("#", '?', 'mailto:')): | |
| continue | |
| if 'wp-content/uploads' in url: | |
| clean_link = url+ "/" + link | |
| else: | |
| clean_link = "https://" + local_domain + "/" + link | |
| if clean_link is not None: | |
| clean_link = clean_link.strip().rstrip('/').replace('/../', '/') | |
| if not any(x in clean_link for x in filter_strings): | |
| clean_links.append(clean_link) | |
| # Return the list of hyperlinks that are within the same domain | |
| return list(set(clean_links)) | |
| # this function will get you a list of all the URLs from the base URL | |
| def crawl(url, local_domain, prog=None): | |
| # Create a queue to store the URLs to crawl | |
| queue = deque([url]) | |
| # Create a set to store the URLs that have already been seen (no duplicates) | |
| seen = set([url]) | |
| # While the queue is not empty, continue crawling | |
| while queue: | |
| # Get the next URL from the queue | |
| url_pop = queue.pop() | |
| # Get the hyperlinks from the URL and add them to the queue | |
| for link in get_domain_hyperlinks(local_domain, url_pop): | |
| if link not in seen: | |
| queue.append(link) | |
| seen.add(link) | |
| if len(seen)>=100: | |
| return seen | |
| if prog is not None: prog(1, desc=f'Crawling: {url_pop}') | |
| return seen | |
| def ingestURL(documents, url, crawling=True, prog=None): | |
| url = url.rstrip('/') | |
| # Parse the URL and get the domain | |
| local_domain = urlparse(url).netloc | |
| if not (local_domain and url.startswith('http')): | |
| return documents | |
| print('Loading URL', url) | |
| if crawling: | |
| # crawl to get other webpages from this URL | |
| if prog is not None: prog(0, desc=f'Crawling: {url}') | |
| links = crawl(url, local_domain, prog) | |
| if prog is not None: prog(1, desc=f'Crawling: {url}') | |
| else: | |
| links = set([url]) | |
| # separate pdf and other links | |
| c_links, pdf_links = [], [] | |
| for x in links: | |
| if x.endswith('.pdf'): | |
| pdf_links.append(x) | |
| elif not x.endswith(media_files): | |
| c_links.append(x) | |
| # Clean links loader using WebBaseLoader | |
| if prog is not None: prog(0.5, desc=f'Ingesting: {url}') | |
| if c_links: | |
| loader = WebBaseLoader(list(c_links)) | |
| documents.extend(loader.load()) | |
| # remote PDFs loader | |
| for pdf_link in list(pdf_links): | |
| loader = PyMuPDFLoader(pdf_link) | |
| doc = loader.load() | |
| for x in doc: | |
| x.metadata['source'] = loader.source | |
| documents.extend(doc) | |
| return documents | |
| def ingestFiles(documents, files_list, prog=None): | |
| for fPath in files_list: | |
| doc = None | |
| if fPath.endswith('.pdf'): | |
| doc = PyMuPDFLoader(fPath).load() | |
| elif fPath.endswith('.txt') and not 'WhatsApp Chat with' in fPath: | |
| doc = TextLoader(fPath).load() | |
| elif fPath.endswith(('.doc', 'docx')): | |
| doc = Docx2txtLoader(fPath).load() | |
| elif 'WhatsApp Chat with' in fPath and fPath.endswith('.csv'): # Convert Whatsapp TXT files to CSV using https://whatstk.streamlit.app/ | |
| doc = WhatsAppChatLoader(fPath).load() | |
| elif fPath.endswith(('.ppt', '.pptx')): | |
| doc = UnstructuredPowerPointLoader(fPath).load() | |
| else: | |
| pass | |
| if doc is not None and doc[0].page_content: | |
| if prog is not None: prog(0.9, desc='Loaded file: '+fPath.rsplit('/')[0]) | |
| print('Loaded file:', fPath) | |
| documents.extend(doc) | |
| return documents | |
| def data_ingestion(inputDir=None, file_list=[], url_list=[], gDriveFolder='', prog=None): | |
| documents = [] | |
| # Ingestion from Google Drive Folder | |
| if gDriveFolder: | |
| opFolder = './tmp/gDriveDocs/' | |
| gdown.download_folder(url=gDriveFolder, output=opFolder, quiet=True) | |
| files = [str(x) for x in Path(opFolder).glob('**/*')] | |
| documents = ingestFiles(documents, files, prog) | |
| # Ingestion from Input Directory | |
| if inputDir is not None: | |
| files = [str(x) for x in Path(inputDir).glob('**/*')] | |
| documents = ingestFiles(documents, files, prog) | |
| if file_list: | |
| documents = ingestFiles(documents, file_list, prog) | |
| # Ingestion from URLs - also try https://python.langchain.com/docs/integrations/document_loaders/recursive_url_loader | |
| if url_list: | |
| for url in url_list: | |
| documents = ingestURL(documents, url, prog=prog) | |
| # Cleanup documents | |
| for x in documents: | |
| if 'WhatsApp Chat with' not in x.metadata['source']: | |
| x.page_content = x.page_content.strip().replace('\n', ' ').replace('\\n', ' ').replace(' ', ' ') | |
| # print(f"Total number of documents: {len(documents)}") | |
| return documents | |
| def split_docs(documents): | |
| # Splitting and Chunks | |
| text_splitter = RecursiveCharacterTextSplitter(chunk_size=2500, chunk_overlap=250) # default chunk size of 4000 makes around 1k tokens per doc. with k=4, this means 4k tokens input to LLM. | |
| docs = text_splitter.split_documents(documents) | |
| return docs | |
| def getSourcesFromMetadata(metadata, sourceOnly=True, sepFileUrl=True): | |
| # metadata: list of metadata dict from all documents | |
| setSrc = set() | |
| for x in metadata: | |
| metadataText = '' # we need to convert each metadata dict into a string format. This string will be added to a set | |
| if x is not None: | |
| # extract source first, and then extract all other items | |
| source = x['source'] | |
| source = source.rsplit('/',1)[-1] if 'http' not in source else source | |
| notSource = [] | |
| for k,v in x.items(): | |
| if v is not None and k!='source' and k in ['page']: | |
| notSource.extend([f"{k}: {v}"]) | |
| metadataText = ', '.join([f'source: {source}'] + notSource) if sourceOnly==False else source | |
| setSrc.add(metadataText) | |
| if sepFileUrl: | |
| src_files = '\n'.join(([f"{i+1}) {x}" for i,x in enumerate(sorted([x for x in setSrc if 'http' not in x], key=str.casefold))])) | |
| src_urls = '\n'.join(([f"{i+1}) {x}" for i,x in enumerate(sorted([x for x in setSrc if 'http' in x], key=str.casefold))])) | |
| src_files = 'Files:\n'+src_files if src_files else '' | |
| src_urls = 'URLs:\n'+src_urls if src_urls else '' | |
| newLineSep = '\n\n' if src_files and src_urls else '' | |
| return src_files + newLineSep + src_urls , len(setSrc) | |
| else: | |
| src_docs = '\n'.join(([f"{i+1}) {x}" for i,x in enumerate(sorted(list(setSrc), key=str.casefold))])) | |
| return src_docs, len(setSrc) | |
| def getEmbeddingFunc(creds): | |
| # OpenAI key used | |
| if creds.get('service')=='openai': | |
| embeddings = OpenAIEmbeddings(openai_api_key=creds.get('oai_key','Null')) | |
| # WX key used | |
| elif creds.get('service')=='watsonx' or creds.get('service')=='bam': | |
| # testModel = Model(model_id=ModelTypes.FLAN_UL2, credentials=creds['credentials'], project_id=creds['project_id']) # test the API key | |
| # del testModel | |
| embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2", cache_folder='./tmp/hfEmbModel') # for now use OpenSource model for embedding as WX doesnt have any embedding model | |
| else: | |
| raise Exception('Error: Invalid or None Credentials') | |
| return embeddings | |
| def getVsDict(embeddingFunc, docs, vsDict={}): | |
| # create chroma client if doesnt exist | |
| if vsDict.get('chromaClient') is None: | |
| vsDict['chromaDir'] = './tmp/vecstore/'+str(uuid.uuid1()) | |
| vsDict['chromaClient'] = Chroma(embedding_function=embeddingFunc, persist_directory=vsDict['chromaDir']) | |
| # clear chroma client before adding new docs | |
| if vsDict['chromaClient']._collection.count()>0: | |
| vsDict['chromaClient'].delete(vsDict['chromaClient'].get()['ids']) | |
| # add new docs to chroma client | |
| vsDict['chromaClient'].add_documents(docs) | |
| print('vectorstore count:',vsDict['chromaClient']._collection.count(), 'at', datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')) | |
| return vsDict | |
| # used for Hardcoded documents only - not uploaded by user (userData_vecStore is separate function) | |
| def localData_vecStore(embKey={}, inputDir=None, file_list=[], url_list=[], vsDict={}, gGrUrl=''): | |
| documents = data_ingestion(inputDir, file_list, url_list, gGrUrl) | |
| if not documents: | |
| raise Exception('Error: No Documents Found') | |
| docs = split_docs(documents) | |
| # Embeddings | |
| embeddings = getEmbeddingFunc(embKey) | |
| # create chroma client if doesnt exist | |
| vsDict_hd = getVsDict(embeddings, docs, vsDict) | |
| # get sources from metadata | |
| src_str = getSourcesFromMetadata(vsDict_hd['chromaClient'].get()['metadatas']) | |
| src_str = str(src_str[1]) + ' source document(s) successfully loaded in vector store.'+'\n\n' + src_str[0] | |
| print(src_str) | |
| return vsDict_hd | |
| def num_tokens_from_string(string, encoding_name = "cl100k_base"): | |
| """Returns the number of tokens in a text string.""" | |
| encoding = tiktoken.get_encoding(encoding_name) | |
| num_tokens = len(encoding.encode(string)) | |
| return num_tokens | |
| def changeModel(oldModel, newModel): | |
| if oldModel: | |
| warning = 'Credentials not found for '+oldModel+'. Using default model '+newModel | |
| gr.Warning(warning) | |
| time.sleep(1) | |
| return newModel | |
| def getModelChoices(openAi_models, wml_models, bam_models): | |
| return [model +' (openai)' for model in openAi_models] + [model.value +' (watsonx)' for model in wml_models] + [model + ' (bam)' for model in bam_models] | |
| def gradio_auth(userInput, passInput): | |
| if userInput == os.getenv("GRADIO_USER",'') and passInput == os.getenv("GRADIO_PASS",''): | |
| return True | |
| else: | |
| return False |