Use transformers

app.py

@@ -7,6 +7,7 @@ from huggingface_hub import HfApi
 from huggingface_hub.utils import logging
 from llama_cpp import Llama
 import pandas as pd
+from transformers import pipeline
 
 load_dotenv()
 
@@ -27,10 +28,13 @@ client = Client(headers=headers)
 api = HfApi(token=HF_TOKEN)
 
 print("About to load DuckDB-NSQL-7B model")
+"""
 llama = Llama(
         model_path="DuckDB-NSQL-7B-v0.1-q8_0.gguf",
         n_ctx=2048,
 )
+"""
+pipe = pipeline("text-generation", model="motherduckdb/DuckDB-NSQL-7B-v0.1")
 print("DuckDB-NSQL-7B model has been loaded")
 
 def get_first_parquet(dataset: str):
@@ -47,6 +51,14 @@ def query_remote_model(text):
     pred = response.json()
     return pred[0]["generated_text"]
 
+def query_local_model_transformers(text):
+    pred = pipe(text)
+    print(type(pred))
+    print(pred)
+    return pred
+    #pred = llama(text, temperature=0.1, max_tokens=500)
+    #return pred["choices"][0]["text"]
+
 
 def query_local_model(text):
     pred = llama(text, temperature=0.1, max_tokens=500)
@@ -84,7 +96,7 @@ def text2sql(dataset_name, query_input):
     
     # sql_output =  query_remote_model(text)
 
-    sql_output = query_local_model(text)
+    sql_output = query_local_model_transformers(text)
 
     try:
         query_result = con.sql(sql_output).df()