Update app.py

app.py

@@ -4,7 +4,6 @@ import gradio as gr
 from dotenv import load_dotenv
 from httpx import Client
 from huggingface_hub import HfApi
-#from llama_cpp import Llama
 import pandas as pd
 from transformers import pipeline
 import spaces
@@ -14,32 +13,14 @@ load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 assert HF_TOKEN is not None, "You need to set HF_TOKEN in your environment variables"
 
-
 BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
-API_URL = "https://m82etjwvhoptr3t5.us-east-1.aws.endpoints.huggingface.cloud"
-headers = {
-	"Accept" : "application/json",
-	"Authorization": f"Bearer {HF_TOKEN}",
-	"Content-Type": "application/json" 
-}
-
 client = Client(headers=headers)
 api = HfApi(token=HF_TOKEN)
-
-# First approach: Use llama.cpp
-#llama = Llama(model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf", n_ctx=2048)
-#def query_local_model(text):
-#    pred = llama(text, temperature=0.1, max_tokens=500)
-#    return pred["choices"][0]["text"]
-
-
-# Second approach: Use transformers -> Took too much time
 pipe = pipeline("text-generation", model="motherduckdb/DuckDB-NSQL-7B-v0.1", device=1)
-# pipe.to('cuda')
 
 @spaces.GPU
-def query_local_model_transformers(prompt):
-    pred = pipe(text, max_length=1000)
+def generate_sql(prompt):
+    pred = pipe(prompt, max_length=1000)
     return pred[0]["generated_text"]
 
 
@@ -48,16 +29,6 @@ def get_first_parquet(dataset: str):
     return resp.json()["parquet_files"][0]
 
 
-def query_remote_model(text):
-    payload = {
-        "inputs": text,
-        "parameters": {}
-    }
-    response = client.post(API_URL, headers=headers, json=payload)
-    pred = response.json()
-    return pred[0]["generated_text"]
-
-
 def text2sql(dataset_name, query_input):
     print(f"start text2sql for {dataset_name}")
     try:
@@ -92,8 +63,7 @@ def text2sql(dataset_name, query_input):
     ### Response (use duckdb shorthand if possible):
     """
     try:
-        #sql_output =  query_remote_model(text)
-        sql_output = query_local_model_transformers(text)
+        sql_output = generate_sql(text)
     except Exception as error:
         return {
             schema_output: ddl_create,
@@ -123,7 +93,7 @@ with gr.Blocks() as demo:
     gr.Markdown("This space showcase how to generate a SQL query from a text and get the result.")
     gr.Markdown("Tech stack: duckdb and DuckDB-NSQL-7B model")
     dataset_name = gr.Textbox("jamescalam/world-cities-geo", label="Dataset Name")
-    query_input = gr.Textbox("Cities from Albania country?", label="Ask something about your data")
+    query_input = gr.Textbox("Cities from Albania country", label="Ask something about your data")
     examples = [
         ["Cities from Albania country"],
         ["The continent with the most number of countries"],