Update app.py

app.py

@@ -7,6 +7,7 @@ from huggingface_hub import HfApi
 import pandas as pd
 from transformers import pipeline
 import spaces
+from llama_cpp import Llama
 
 load_dotenv()
 
@@ -22,12 +23,16 @@ headers = {
 client = Client(headers=headers)
 api = HfApi(token=HF_TOKEN)
 pipe = pipeline("text-generation", model="motherduckdb/DuckDB-NSQL-7B-v0.1", device="cuda")
-
+llama = Llama(
+        model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf",
+        n_ctx=2048,
+    )
 @spaces.GPU
 def generate_sql(prompt):
-    pred = pipe(prompt, max_length=1000)
-    return pred[0]["generated_text"]
-
+    # pred = pipe(prompt, max_length=1000)
+    # return pred[0]["generated_text"]
+    pred = llama(prompt, temperature=0.1, max_tokens=1000)
+    return pred["choices"][0]["text"]
 
 def get_first_parquet(dataset: str):
     resp = client.get(f"{BASE_DATASETS_SERVER_URL}/parquet?dataset={dataset}")