借助 NVIDIA NIM 加速 Vanna 上的文本到 SQL 推理，从而提升分析效率

from pymilvus import MilvusClient, model
from vanna.milvus import Milvus_VectorStore
from vanna.openai import OpenAI_Chat
from openai import OpenAI

# Set the nvidia api key 
nvidia_api_key = '...'

# Define Vanna Class
class VannaMilvus(Milvus_VectorStore, OpenAI_Chat):
    def __init__(self, llm_client, config=None):
        Milvus_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, client=llm_client, config=config)

def get_openai_client():
    client = OpenAI(
        base_url = "https://integrate.api.nvidia.com/v1",
        api_key = nvidia_api_key
    )
    return client

llm_client = get_openai_client()

class EmbeddingWrapper:
    def __init__(self, embedder):
        self.embedder = embedder

    def encode_documents(self, texts):
        result = self.embedder.embed_documents(texts)
        return [np.array(r) for r in result]

    def encode_queries(self, texts):
        embeddings = []
        for text in texts:
            embeddings.append(self.embedder.embed_query(text))
        return embeddings

vanna_embedder = EmbeddingWrapper(nvidia_embedder)

def get_openai_client():
    client = OpenAI(
        base_url = "https://integrate.api.nvidia.com/v1",
        api_key = nvidia_api_key
    )
    return client

llm_client = get_openai_client()

milvus_uri = "./milvus_nvidia.db"
milvus_client_nvidia = MilvusClient(uri=milvus_uri)

# Specify the LLM used for SQL generation
model_name = "nvidia/llama-3.1-70b-instruct"

# Define the configuration for Vanna instance
config_nvidia = {
        "model": model_name,
        "milvus_client": milvus_client_nvidia,
        "embedding_function": vanna_embedder,
        "n_results": 2,  # The number of results to return from Milvus semantic search.
    }

vn_nvidia = VannaMilvus(llm_client, config=config_nvidia)

import sqlite3
import pandas as pd

# Specify the path to the SQLite database
sqlite_path = 'steam_data.db'

# Connect to the SQLite database
sql_connect = sqlite3.connect(sqlite_path)
c = sql_connect.cursor()

# Create tables
init_sqls = """
CREATE TABLE IF NOT EXISTS games (
    app_id INTEGER PRIMARY KEY,
    name TEXT,
    release_date TEXT,
    price REAL,
    short_description TEXT,
    positive INTEGER,
    negative INTEGER,
    min_owners INTEGER,
    max_owners INTEGER,
    hltb_single REAL
);

CREATE TABLE IF NOT EXISTS categories (
    app_id INTEGER,
    categories TEXT,
    FOREIGN KEY (app_id) REFERENCES games(app_id)
);

CREATE TABLE IF NOT EXISTS tags (
    app_id INTEGER,
    tags TEXT,
    tag_frequencies TEXT,
    FOREIGN KEY (app_id) REFERENCES games(app_id)
);
"""

for sql in init_sqls.split(";"):
    c.execute(sql)

# Read the CSV files
games_df = pd.read_csv('processed_dataset/games.csv')
categories_df = pd.read_csv('processed_dataset/categories.csv')
tags_df = pd.read_csv('processed_dataset/tags.csv')

# Insert data into tables
games_df.to_sql('games', sql_connect, if_exists='append', index=False)
categories_df.to_sql('categories', sql_connect, if_exists='append', index=False)
tags_df.to_sql('tags', sql_connect, if_exists='append', index=False)

sql_connect.commit()

# Connect to the SQLite database
vn_nvidia.connect_to_sqlite(sqlite_path)

# Remove existing training data
existing_training_data = vn_nvidia.get_training_data()
if len(existing_training_data) > 0:
    for _, training_data in existing_training_data.iterrows():
        vn_nvidia.remove_training_data(training_data["id"])

# Get the DDL of the SQLite database
df_ddl = vn_nvidia.run_sql("SELECT type, sql FROM sqlite_master WHERE sql is not null")

# Train the model on the DDL data
for ddl in df_ddl["sql"].to_list():
    vn_nvidia.train(ddl=ddl)

# Add documentation about your business terminology or definitions.
vn_nvidia.train(
    documentation="""
    This dataset is used to answer questions about the game trends.
    """
)

# Add documentation about the tables
vn_nvidia.train(
    documentation="""
    The games table contains information about the games.
    The app_id is the unique identifier for the game.This is a primary key.
    The name is the name of the game.
    The release_date is the date the game was released.
    The price is the price of the game. Price in USD, 0.0 if its free.
    The short_description is a brief description of the game.
    The positive is the number of positive reviews or votes.
    The negative is the number of negative reviews or votes.
    The min_owners is the minimum number of owners. Used together with max_owners to get an estimate of the player base.
    The max_owners is the maximum number of owners. Used together with min_owners to get an estimate of the player base.
    The hltb_single is the average playtime of the game. This is an estimate.
    """
)

vn_nvidia.train(
    documentation="""
    The categories table contains information about the categories of the games.
    The app_id is the unique identifier for the game.
    The categories is the categories of the game.
    The app_id is a foreign key to the games table.
    """
)

vn_nvidia.train(
    documentation="""
    The tags table contains information about the tags of the games.
    The app_id is the unique identifier for the game.
    The tags is the tags of the game. These are user defined.
    The tag_frequencies is the frequencies of the tags.
    The app_id is a foreign key to the games table.
    """
)

training_data = vn_nvidia.get_training_data()
training_data

sql = vn_nvidia.generate_sql("Which 5 games have the most positive reviews and how many?")
vn_nvidia.run_sql(sql)

sql = vn_nvidia.generate_sql("Which indie game has the biggest player base?")
vn_nvidia.run_sql(sql)