make model parameters more dynamic w env variables

Hermes-2-Pro-Llama-3-8B-Q8_0.gguf

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d138388cfda04d185a68eaf2396cf7a5cfa87d038a20896817a9b7cf1806f532
-size 8541050176

app.py

@@ -1,5 +1,6 @@
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from llama_cpp.llama_speculative import LlamaPromptLookupDecoding
+from huggingface_hub import hf_hub_download
 from huggingface_hub import HfApi
 import matplotlib.pyplot as plt
 from typing import Tuple, Optional
@@ -11,6 +12,7 @@ import llama_cpp
 import instructor
 import spaces
 import enum
+import os
 
 from pydantic import BaseModel, Field
 
@@ -20,6 +22,18 @@ view_name = "dataset_view"
 hf_api = HfApi()
 conn = duckdb.connect()
 
+gpu_layers = int(os.environ.get("GPU_LAYERS", 81))
+draft_pred_tokens = int(os.environ.get("DRAFT_PRED_TOKENS", 2))
+
+repo_id = os.getenv("MODEL_REPO_ID", "NousResearch/Hermes-2-Pro-Llama-3-8B-GGUF")
+model_file_name = os.getenv("MODEL_FILE_NAME", "Hermes-2-Pro-Llama-3-8B-Q8_0.gguf")
+
+hf_hub_download(
+    repo_id=repo_id,
+    filename=model_file_name,
+    local_dir="./models",
+)
+
 
 class OutputTypes(str, enum.Enum):
     TABLE = "table"
@@ -75,10 +89,10 @@ CREATE TABLE {} (
 @spaces.GPU(duration=120)
 def generate_query(ddl: str, query: str) -> dict:
     llama = llama_cpp.Llama(
-        model_path="Hermes-2-Pro-Llama-3-8B-Q8_0.gguf",
-        n_gpu_layers=50,
+        model_path=f"models/{model_file_name}",
+        n_gpu_layers=gpu_layers,
         chat_format="chatml",
-        draft_model=LlamaPromptLookupDecoding(num_pred_tokens=2),
+        draft_model=LlamaPromptLookupDecoding(num_pred_tokens=draft_pred_tokens),
         logits_all=True,
         n_ctx=2048,
         verbose=True,
@@ -94,16 +108,13 @@ def generate_query(ddl: str, query: str) -> dict:
     You are an expert SQL assistant with access to the following PostgreSQL Table:
     
     ```sql
-    {ddl}
+    {ddl.strip()}
     ```
     
     Please assist the user by writing a SQL query that answers the user's question.
-    
-    Use Label Key as the column name for the x-axis and Data Key as the column name for the y-axis for chart responses. The 
-    label key and data key must be present in the SQL output.
     """
 
-    print("Calling LLM with system prompt: ", system_prompt)
+    print("Calling LLM with system prompt: ", system_prompt, query)
 
     resp: SQLResponse = create(
         model="Hermes-2-Pro-Llama-3-8B",
@@ -135,6 +146,7 @@ def query_dataset(dataset_id: str, query: str) -> Tuple[pd.DataFrame, str, plt.F
     data_key = response.get("data_key")
     viz_type = response.get("visualization_type")
     sql = response.get("sql")
+    markdown_output = f"""```sql\n{sql}\n```"""
 
     # handle incorrect data and label keys
     if label_key and label_key not in df.columns:
@@ -142,6 +154,9 @@ def query_dataset(dataset_id: str, query: str) -> Tuple[pd.DataFrame, str, plt.F
     if data_key and data_key not in df.columns:
         data_key = None
 
+    if df.empty:
+        return df, f"```sql\n{sql}\n```", plot
+
     if viz_type == OutputTypes.LINECHART:
         plot = df.plot(kind="line", x=label_key, y=data_key).get_figure()
         plt.xticks(rotation=45, ha="right")
@@ -151,7 +166,6 @@ def query_dataset(dataset_id: str, query: str) -> Tuple[pd.DataFrame, str, plt.F
         plt.xticks(rotation=45, ha="right")
         plt.tight_layout()
 
-    markdown_output = f"""```sql\n{sql}\n```"""
     return df, markdown_output, plot
 
 
@@ -167,8 +181,8 @@ with gr.Blocks() as demo:
     examples = [
         ["Show me a preview of the data"],
         ["Show me something interesting"],
-        ["What is the largest length of sql query context?"],
-        ["show me counts by sql_query_type in a bar chart"],
+        ["Which row has longest description length?"],
+        ["find the average length of sql query context"],
     ]
     gr.Examples(examples=examples, inputs=[user_query], outputs=[])