import os import gradio import shutil import json import dotenv import pickle import itertools from collections import Counter from datetime import datetime from huggingface_hub import HfApi, HfFolder from openai import OpenAI as OpenAIClient from pinecone import ServerlessSpec from pinecone.grpc import PineconeGRPC as Pinecone from transformers import BertTokenizerFast from langchain.text_splitter import CharacterTextSplitter from langchain_community.document_loaders import DirectoryLoader ''' This file contains the code for user prompting of the language model. The language model used is gpt 3.5 turbo and uses documents stored in Pinecone. ''' # Load environment variables dotenv.load_dotenv() assert os.getenv("OPENAI_API_KEY") is not None, "Please set the OPENAI_API_KEY environment variable." assert os.getenv("PINECONE_API_KEY") is not None, "Please set the PINECONE_API_KEY environment variable." assert os.getenv("HUGGINGFACE_API_KEY") is not None, "Please set the HUGGINGFACE_API_KEY environment variable." HfFolder.save_token(os.getenv("HUGGINGFACE_API_KEY")) CHAT_HISTORY_FILE = "chat_history.pkl" NEW_UPLOAD_DIRECTORY = "new_uploads/" PREV_UPLOAD_DIRECTORY = "prev_uploads/" DOC_CHUNK_SIZE = 1000 DOC_CHUNK_OVERLAP = 40 EMBEDDING_FILE = 'embeddings.json' INDEX_NAME = "hybrid" BATCH_SIZE = 100 INTERACTIONS_DATASET = "ryanRocks/FalconOpenAIInteractions" FILES_DATASET = "ryanRocks/FalconOpenAIFiles" # Ensure uploads directory exists os.makedirs(NEW_UPLOAD_DIRECTORY, exist_ok=True) os.makedirs(PREV_UPLOAD_DIRECTORY, exist_ok=True) chat_history = [] chatbot_history = [] # Initialize Pinecone database try: pc = Pinecone( api_key=os.getenv("PINECONE_API_KEY"), pool_threads=30, spec=ServerlessSpec( cloud="aws", region="us-east-1", ), ) print("Connected to Pinecone") except Exception as e: print(f"Error initializing Pinecone: {e}") exit() # Initialize index if it does not exist existing_indexes = [index.name for index in pc.list_indexes().indexes] if INDEX_NAME not in existing_indexes: pc.create_index( name=INDEX_NAME, dimension=1536, metric="dotproduct", spec=ServerlessSpec( cloud="aws", region="us-east-1", ), ) print(f"Created index {INDEX_NAME}") else: print(f"Index {INDEX_NAME} already exists") def initialize_chat_history(): ''' Initialize the chat history using the chat history pickle file. ''' global chatbot_history loaded_chat_history = [] if (os.path.exists(CHAT_HISTORY_FILE)): with open(CHAT_HISTORY_FILE, "rb") as f: loaded_chat_history = pickle.load(f) else: loaded_chat_history = [] for i in range(0, len(loaded_chat_history), 2): chatbot_history.append(( loaded_chat_history[i]['content'], loaded_chat_history[i+1]['content'] )) async def upload_file(files): ''' Upload files to Pinecone Args: files: List of file paths to process ''' # Copy files to uploads directory for easier processing for file in files: file_path = os.path.join(NEW_UPLOAD_DIRECTORY, file.name.split('/')[-1]) shutil.move(file.name, file_path) # Load documents documents = read_documents() dense_embeddings = dense_embed(documents) sparse_embeddings = sparse_embed(documents) #save_embeddings(dense_embeddings, EMBEDDING_FILE) # Upsert embeddings into Pinecone await upsert(dense_embeddings, sparse_embeddings) # Move newly uploaded files to previous uploads directory move_files() return get_uploaded_files() def read_documents(): ''' Load documents from a specified directory into a list Args: file_paths: List of file paths to load documents from Returns: documents: List of documents loaded from the directory, split by chunks ''' # Load documents print("Loading documents...") documents = [] # Declare loaders for different file types pdf_loader = DirectoryLoader(NEW_UPLOAD_DIRECTORY, glob="*.pdf") docx_loader = DirectoryLoader(NEW_UPLOAD_DIRECTORY, glob="*.docx") txt_loader = DirectoryLoader(NEW_UPLOAD_DIRECTORY, glob="*.txt") for loader in [pdf_loader, docx_loader, txt_loader]: # Load document try: # Error loading documents: Expected directory, got file: '/private/var/folders/sc/_5mj781j5315nzv10s8kvs1w0000gn/T/gradio/33a9766ee3f05c368d3c7fe56f6f2356e88a4348/YuYouChen Resume.pdf' documents.extend(loader.load()) except Exception as e: print(f"Error loading documents: {e}") if (len(documents) == 0): print("No documents loaded.") return [] # Split documents into chunks text_splitter = CharacterTextSplitter(chunk_size=DOC_CHUNK_SIZE, chunk_overlap=DOC_CHUNK_OVERLAP) documents = text_splitter.split_documents(documents) # Iterate to edit metadata to include chunk number # format = {filename}_{chunk number} chunk_num = 1 prev_doc_id = documents[0].metadata['source'] for chunk in documents: if chunk.metadata['source'] != prev_doc_id: chunk_num = 1 prev_doc_id = chunk.metadata['source'] chunk.metadata['source'] = f"{prev_doc_id}_{chunk_num}" chunk_num += 1 print("Documents loaded") return documents def dense_embed(documents): ''' Embed documents using OpenAIEmbeddings Args: documents: List of documents to embed Returns: List of JSON objects {doc_id, embeddings, metadata} ''' print("Generating dense embeddings...") # Use OpenAI to embed documents client = OpenAIClient( api_key=os.getenv("OPENAI_API_KEY") ) embeddings = [] # Embed each chunk for chunk in documents: chunk_embeddings = client.embeddings.create( model="text-embedding-3-small", input=chunk.page_content ) # Extract embeddings from response chunk_embedding = [record.embedding for record in chunk_embeddings.data] embeddings.append({ 'doc_id': chunk.metadata['source'].split('/')[-1], 'embeddings': chunk_embedding[0], 'metadata': {'source': chunk.metadata['source'], 'text': chunk.page_content} }) print("Complete") return embeddings def sparse_embed(documents): ''' Generate sparse embeddings for a list of documents Args: documents: List of documents to generate sparse embeddings for Returns: List of sparse embeddings in dictionary format ''' print("Generating sparse embeddings...") tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased') sparse_embeds = [] for chunk in documents: # Create batch of input_ids inputs = tokenizer( chunk.page_content, padding=True, truncation=True, max_length=512, add_special_tokens=False, )['input_ids'] # Create sparse dictionaries sparse_embed = build_dict(inputs) sparse_embeds.append(sparse_embed) print("Complete") return sparse_embeds def build_dict(input_batch): ''' Build a dictionary for sparse embeddings Args: input_batch: List of embeddings to convert to a dictionary Returns: List of sparse embeddings in dictionary format ''' sparse_emb = [] # Iterate through input batch indices = [] values = [] # Convert the input_batch list to a dictionary of key to frequency values freqs = dict(Counter(input_batch)) for idx in freqs: indices.append(idx) values.append(float(freqs[idx])) sparse_emb.append({'indices': indices, 'values': values}) return sparse_emb def save_embeddings(embeddings, filename): ''' Save generated embedding to a json file Args: embeddings: List of embeddings to save filename: Name of the file to save the embeddings ''' print("Saving embedding...") with open(filename, 'w') as file: json.dump(embeddings, file) print("Complete") def chunks(iterable): ''' Breaks vector list into chunks of BATCH_SIZE for parallel upserts Args: iterable: List of vectors to chunk ''' print("Chunking") it = iter(iterable) chunk = tuple(itertools.islice(it, BATCH_SIZE)) while chunk: yield chunk chunk = tuple(itertools.islice(it, BATCH_SIZE)) print("Complete") def vectorize(dense_embeddings, sparse_embeddings): ''' Vectorize embeddings with document ids to prepare for insertion into Pinecone Args: embeddings: List of embeddings to vectorize Returns: List of vectors with tuples (chunk ids, embeddings) ''' print("Vectorizing...") vectors = [] for dense, sparse in zip(dense_embeddings, sparse_embeddings): vectors.append({ 'id': dense['doc_id'], 'values': dense['embeddings'], 'sparse_values': sparse[0], 'metadata': dense['metadata'], }) print("Vectorized") return vectors async def upsert(dense_embeddings, sparse_embeddings): ''' Upsert embeddings into pinecone ''' index = pc.Index(INDEX_NAME) vectors = vectorize(dense_embeddings, sparse_embeddings) # Insert vectors into database in chunks print("Upserting embeddings...") vector_chunks = chunks(vectors) for chunk in vector_chunks: index.upsert(chunk) print("Complete") def move_files(): ''' Move uploaded files to the previous uploads directory ''' print("Moving files...") api = HfApi() for file in os.listdir(NEW_UPLOAD_DIRECTORY): file_path = os.path.join(NEW_UPLOAD_DIRECTORY, file) if (os.path.isfile(file_path)): # Upload file to HuggingFace Datasets api.upload_file( path_or_fileobj = file_path, path_in_repo = file, repo_id = "ryanRocks/FalconOpenAIFiles", repo_type = "dataset", ) # Move file to previous uploads directory new_file_path = os.path.join(PREV_UPLOAD_DIRECTORY, file) shutil.move(file_path, new_file_path) def hybrid_scale(dense, sparse, alpha): print("Hybrid scaling...") # Check alpha value in range 0 to 1 if alpha < 0 or alpha > 1: raise ValueError("Alpha must be between 0 and 1") # Scale dense and sparse vectors to create hybrid search vectors hdense = [v * alpha for v in dense] hsparse = { 'indices': sparse['indices'], 'values': [v * (1 - alpha) for v in sparse['values']], } print("Complete") return hdense, hsparse def hybrid_query(question, top_k, alpha): try: print("Hybrid querying...") # Convert the question into a dense vector print("Converting question to dense vector...") client = OpenAIClient( api_key=os.getenv("OPENAI_API_KEY") ) query_embedding = client.embeddings.create( model="text-embedding-3-small", input=question, ) dense_vec = [record.embedding for record in query_embedding.data][0] print("Complete") # Convert the question into a sparse vector print("Converting question to sparse vector...") tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased') inputs = tokenizer( question, padding=True, truncation=True, max_length=512, add_special_tokens=False, )['input_ids'] sparse_vec = build_dict(inputs)[0] print("Complete") # Scale alpha with hybrid_scale dense_vec, sparse_vec = hybrid_scale( dense_vec, sparse_vec, alpha ) # Query pinecone with the query parameters print("Querying Pinecone...") index = pc.Index(INDEX_NAME) result = index.query( vector=dense_vec, sparse_vector=sparse_vec, top_k=top_k, include_values=True, include_metadata=True, ) print("Complete") # Return search results as json return result except Exception as e: print(f"Error querying Pinecone: {e}") def prompt(question, history): ''' Prompt the language model with user input Args: question: User string input to prompt the language model Returns: Language model response to the user ''' global chatbot_history global chat_history # Handle clearing history if len(history) == 0: chatbot_history = [] chat_history = [] with open(CHAT_HISTORY_FILE, "wb") as f: pickle.dump(chat_history, f) # Check for API key if (os.getenv("OPENAI_API_KEY") is None): return "Please set the OPENAI_API_KEY environment variable." # Query database and prompt the language model with results try: # Query database for context results = hybrid_query(question, top_k=5, alpha=0.4) context = "" for match in results['matches']: context += match['metadata']['text'] + "\n" client = OpenAIClient() # Prepare prompt with chat history prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:" messages = [{"role": "system", "content": "You are a helpful assistant."}] messages.extend(chat_history) messages.append({"role": "user", "content": prompt}) time = datetime.now().isoformat() print("Prompting language model...") response = client.chat.completions.create( messages=messages, model="gpt-3.5-turbo", ) # Extract answer from response answer = response.choices[0].message.content.strip() interaction = [ {"timestamp": time}, {"role": "user", "content": question}, {"role": "assistant", "content": answer}, {"full_prompt": messages}, ] with open(f"{time}.json", "w") as f: json.dump(interaction, f) # Upload interaction to HuggingFace Datasets api = HfApi() api.upload_file( path_or_fileobj = f"{time}.json", path_in_repo = f"{time}.json", repo_id = INTERACTIONS_DATASET, repo_type = "dataset", ) # Save chat history chat_history.extend([ {"role": "user", "content": question}, {"role": "assistant", "content": answer}, ]) with open(CHAT_HISTORY_FILE, "wb") as f: pickle.dump(chat_history, f) return answer # Handle exceptions except Exception as e: return "Error: " + str(e) def get_uploaded_files(): ''' Get uploaded files in prev_uploads directory ''' uploaded_files = [] for file in os.listdir(PREV_UPLOAD_DIRECTORY): file_path = os.path.join(PREV_UPLOAD_DIRECTORY, file) if (os.path.isfile(file_path)): uploaded_files.append(file_path) return uploaded_files # Create a Gradio interface with gradio.Blocks() as demo: # Load chat history initialize_chat_history() # Create chatbot interface chatbot = gradio.Chatbot(value=chat_history, placeholder="What would you like to know?") gradio.ChatInterface(fn=prompt, chatbot=chatbot) # Create file upload interface file_output = gradio.File(value=get_uploaded_files()) upload_button = gradio.UploadButton("Click to upload a file", file_types=["pdf, docx, txt"], file_count="multiple") upload_button.upload(upload_file, upload_button, file_output) demo.launch()