Update app.py

Some Optimization.

app.py

@@ -3,18 +3,23 @@ import gradio as gr
 from transformers import pipeline
 import torch
 
+DEVICE = 0 if torch.cuda.is_available() else -1
+
+generator = pipeline(
+    "text-generation",
+    model="sirunchained/text-to-sql-model-v2",
+    device=DEVICE,
+    torch_dtype=torch.float16 if DEVICE == 0 else torch.float32,
+)
 
 @spaces.GPU
+@torch.inference_mode()
 def predict_sql(text_input, schema_input=""):
-    # Get model
-    DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    generator = pipeline("text-generation", model="sirunchained/text-to-sql-model-v2", device=DEVICE)
-    
     # Construct the user content including schema if provided
-    if schema_input and schema_input.strip() != "":
-        user_content = f"# Schema \n{schema_input} \n\n# Text \n{text_input}"
+    if schema_input and schema_input.strip():
+        user_content = f"# Schema\n{schema_input}\n\n# Text\n{text_input}"
     else:
-        user_content = f"# Text \n{text_input}"
+        user_content = f"# Text\n{text_input}"
 
     chat_template_input = [
         {"role": "user", "content": user_content}
@@ -26,21 +31,23 @@ def predict_sql(text_input, schema_input=""):
         tokenize=False,
         add_generation_prompt=True,
     )
+
     # Generate output using the fine-tuned model
-    output = generator(formatted_input, max_new_tokens=256, disable_compile=True)
+    output = generator(
+        formatted_input,
+        max_new_tokens=256,
+        do_sample=False,
+    )
 
     # Extract the generated SQL query
     generated_text = output[0]["generated_text"]
-    # The model output includes the prompt, so we need to extract only the assistant's part
-    # The assistant's response starts after '<start_of_turn>model'
-    start_token = '<start_of_turn>model\n'
-    if start_token in generated_text:
-        sql_query = generated_text.split(start_token)[-1].strip()
-    else:
-        sql_query = generated_text.strip() # Fallback if token not found
+
+    # Remove the prompt from the generated output
+    sql_query = generated_text[len(formatted_input):].strip()
 
     # Further clean up any unwanted tokens like '<end_of_turn>'
-    sql_query = sql_query.replace('<end_of_turn>', '').strip()
+    sql_query = sql_query.replace("<end_of_turn>", "").strip()
+
     return sql_query
 
 # Sample queries for the Gradio interface including SELECT, INSERT, UPDATE, and DELETE operations