Added Langsmith

app.py

@@ -4,8 +4,9 @@ import duckdb
 import spaces
 import gradio as gr
 import pandas as pd
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
+from langchain_huggingface.llms import HuggingFacePipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+from langsmith import traceable
 
 
 # Height of the Tabs Text Area
@@ -36,6 +37,9 @@ quantization_config = BitsAndBytesConfig(
 
 model = AutoModelForCausalLM.from_pretrained("motherduckdb/DuckDB-NSQL-7B-v0.1", quantization_config=quantization_config,
                                              device_map="auto", torch_dtype=torch.bfloat16)
+
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024, return_full_text=False)
+hf = HuggingFacePipeline(pipeline=pipe)
 print('Model Loaded...')
 print(f'Model Device: {model.device}')
 
@@ -88,13 +92,10 @@ def get_prompt(schema, query_input):
     return text
 
 @spaces.GPU(duration=60)
+@traceable()
 def generate_sql(prompt):
-
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-    input_token_len = input_ids.shape[1]
-    outputs = model.generate(input_ids.to(model.device), max_new_tokens=1024)
-    result = tokenizer.decode(outputs[0][input_token_len:], skip_special_tokens=True)
-    return result
+    result = hf.invoke(prompt)
+    return result.strip()
 
 # Generate SQL
 def text2sql(table, query_input):

requirements.txt

@@ -1,4 +1,6 @@
 accelerate
 bitsandbytes
 transformers
-duckdb
+duckdb
+langsmith
+langchain-huggingface