Update app.py

app.py

@@ -6,7 +6,8 @@ from httpx import Client
 from huggingface_hub import HfApi
 #from llama_cpp import Llama
 import pandas as pd
-#from transformers import pipeline
+from transformers import pipeline
+import spaces
 
 load_dotenv()
 
@@ -33,10 +34,13 @@ api = HfApi(token=HF_TOKEN)
 
 
 # Second approach: Use transformers -> Took too much time
-#pipe = pipeline("text-generation", model="motherduckdb/DuckDB-NSQL-7B-v0.1")
-#def query_local_model_transformers(text):
-#    pred = pipe(text, max_length=1000)
-#    return pred[0]["generated_text"]
+pipe = pipeline("text-generation", model="motherduckdb/DuckDB-NSQL-7B-v0.1")
+pipe.to('cuda')
+
+@spaces.GPU
+def query_local_model_transformers(prompt):
+    pred = pipe(text, max_length=1000)
+    return pred[0]["generated_text"]
 
 
 def get_first_parquet(dataset: str):