chatbot/app2.pyn

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