import os import gradio as gr import pandas as pd import numpy as np import chromadb from chromadb.config import Settings from io import StringIO from sentence_transformers import SentenceTransformer import plotly.express as px from sklearn.manifold import TSNE # Constants for Model Configuration MODEL_CONFIG = { "gpt-4": { "endpoint": "https://roger-m38jr9pd-eastus2.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2024-08-01-preview", "api_key": os.getenv("GPT4_API_KEY") }, "gpt-4o": { "endpoint": "https://roger-m38jr9pd-eastus2.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-08-01-preview", "api_key": os.getenv("GPT4O_API_KEY") }, "gpt-35-turbo": { "endpoint": "https://rogerkoranteng.openai.azure.com/openai/deployments/gpt-35-turbo/chat/completions?api-version=2024-08-01-preview", "api_key": os.getenv("GPT35_TURBO_API_KEY") }, "gpt-4-32k": { "endpoint": "https://roger-m38orjxq-australiaeast.openai.azure.com/openai/deployments/gpt-4-32k/chat/completions?api-version=2024-08-01-preview", "api_key": os.getenv("GPT4_32K_API_KEY") } } # Initialize Chroma client with DuckDB and Parquet for persistence chroma_client = chromadb.Client() # Functions for Data Processing and Embedding def process_csv_text(temp_file): """Process the uploaded CSV and return the dataframe and column options.""" if isinstance(temp_file, str): df = pd.read_csv(StringIO(temp_file)) else: df = pd.read_csv(temp_file.name, header='infer', sep=',') return df, gr.Dropdown.update(choices=list(df.columns)) def insert_or_update_chroma(col, table, model_name, similarity_metric): """Insert or update embeddings in ChromaDB.""" try: collection = chroma_client.create_collection( name="my_collection", embedding_function=SentenceTransformer(model_name), metadata={"hnsw:space": similarity_metric} ) except Exception: print("Collection exists, deleting it") chroma_client.delete_collection(name='my_collection') collection = chroma_client.create_collection( name="my_collection", embedding_function=SentenceTransformer(model_name), metadata={"hnsw:space": similarity_metric} ) if collection: try: collection.add( documents=list(table[col]), metadatas=[{"source": i} for i in range(len(table))], ids=[str(i + 1) for i in range(len(table))] ) return "Embedding calculations and insertions successful" except Exception as e: return f"Error in embedding calculations: {e}" def show_fig(): """Show t-SNE 2D plot for embeddings.""" collection = chroma_client.get_collection(name="my_collection") embeddings = collection.get(include=['embeddings', 'documents']) df = pd.DataFrame({ 'text': embeddings['documents'], 'embedding': embeddings['embeddings'] }) embeddings_np = np.array(df['embedding'].tolist()) tsne = TSNE(n_components=2, random_state=42) transformed = tsne.fit_transform(embeddings_np) df['tsne_x'] = transformed[:, 0] df['tsne_y'] = transformed[:, 1] fig = px.scatter(df, x='tsne_x', y='tsne_y', hover_name='text') return fig, transformed def show_test_string_fig(test_string, tsne, model_name, similarity_metric): """Show t-SNE plot with test string to compare embeddings.""" collection = chroma_client.get_collection(name="my_collection", embedding_function=SentenceTransformer(model_name)) collection.add( documents=[test_string], metadatas=[{"source": 'test'}], ids=['test_sample'] ) embeddings = collection.get(include=['embeddings', 'documents']) df = pd.DataFrame({ 'text': embeddings['documents'], 'embedding': embeddings['embeddings'], 'set': ['orig' if document != test_string else 'test_string' for document in embeddings["documents"]] }) embeddings_np = np.array(df['embedding'].tolist()) transformed = tsne.transform(embeddings_np) df['tsne_x'] = transformed[:, 0] df['tsne_y'] = transformed[:, 1] fig = px.scatter(df, x='tsne_x', y='tsne_y', hover_name='text', color='set') return fig, tsne def ask_gpt(message, messages_history, embedding_model, system_prompt, temperature, max_tokens, chatgpt_model): """Interacts with the OpenAI API using Azure endpoint.""" if len(messages_history) < 1: messages_history = [{"role": "system", "content": system_prompt}] model_info = MODEL_CONFIG[chatgpt_model] headers = {"Content-Type": "application/json", "api-key": model_info["api_key"]} message = retrieve_similar(message, embedding_model) messages_history += [{"role": "user", "content": message}] response = openai.ChatCompletion.create( model=chatgpt_model, messages=messages_history, temperature=temperature, max_tokens=max_tokens ) return response['choices'][0]['message']['content'], messages_history def retrieve_similar(prompt, embedding_model): """Retrieve similar documents from ChromaDB to enhance context.""" # Initialize SentenceTransformer correctly embedding_function = SentenceTransformer(embedding_model) collection = chroma_client.get_collection( name="my_collection", embedding_function=embedding_function ) results = collection.query(query_texts=prompt, n_results=10) additional_context = '' for i, document in enumerate(results['documents'][0]): additional_context += f'{i + 1}. {document}\n' return additional_context + f'Question: {prompt}' # Gradio Interface Setup def build_gradio_ui(): """Setup the complete Gradio UI.""" with gr.Blocks() as demo: # Tab 1: Upload CSV and Display Data with gr.Tab("Upload data"): upload_button = gr.File(label="Upload CSV", file_types=['.csv'], file_count="single") table = gr.Dataframe(type="pandas", interactive=True) cols = gr.Dropdown(choices=[], label='Dataframe columns') upload_button.change(fn=process_csv_text, inputs=upload_button, outputs=[table, cols]) # Tab 2: ChromaDB, Embeddings, and Plotting with gr.Tab("ChromaDB and Embeddings"): cols = gr.Dropdown(choices=[], label='Dataframe columns') embedding_model = gr.Dropdown(value='all-MiniLM-L6-v2', choices=['all-MiniLM-L6-v2', 'intfloat/e5-small-v2', 'intfloat/e5-base-v2', 'intfloat/e5-large-v2', 'paraphrase-multilingual-MiniLM-L12-v2'], label='Embedding Model') similarity_metric = gr.Dropdown(value='cosine', choices=['cosine', 'l2'], label='Similarity Metric') embedding_button = gr.Button(value="Insert or Update Embeddings") text = gr.Textbox(label='Process Status') show_embeddings_button = gr.Button(value="Show Embeddings") embeddings_plot = gr.Plot() tsne = gr.State(value=None) # Using gr.State for intermediate results (tsne) test_string = gr.Textbox(label='Test String') calculate_2d_repr_button = gr.Button(value="Calculate 2D Representation") embeddings_plot_with_text_string = gr.Plot() embedding_button.click(insert_or_update_chroma, inputs=[cols, table, embedding_model, similarity_metric], outputs=[text]) show_embeddings_button.click(show_fig, inputs=[], outputs=[embeddings_plot, tsne]) calculate_2d_repr_button.click(show_test_string_fig, inputs=[test_string, tsne, embedding_model, similarity_metric], outputs=[embeddings_plot_with_text_string, tsne]) # Tab 3: Chat with GPT with gr.Tab("Chat"): system_prompt = gr.Textbox(value="You are a helpful assistant.", label="System Message") chatgpt_model = gr.Dropdown(value="gpt-4", choices=list(MODEL_CONFIG.keys()), label="ChatGPT Model") temperature = gr.Slider(minimum=0, maximum=2, step=0.1, value=0.7, label="Temperature") max_tokens = gr.Slider(minimum=50, maximum=2000, step=50, value=300, label="Max Tokens") chatbot = gr.Chatbot(label="ChatGPT Chat") clear_button = gr.Button("Clear Chat History") msg = gr.Textbox() msg_log = gr.Textbox("Message history", label="History") # Replacing `.submit()` with `.change()` to trigger callback when user enters a message msg.submit(fn=ask_gpt, inputs=[msg, chatbot, system_prompt, embedding_model, temperature, max_tokens, chatgpt_model], outputs=[msg, chatbot]) clear_button.click(fn=lambda: None, inputs=None, outputs=[chatbot]) return demo # Launch the Gradio interface demo = build_gradio_ui() demo.launch(server_name="0.0.0.0", server_port=8080, share=True) # Launch the Gradio interf