Update app.py

app.py

@@ -51,6 +51,8 @@ def generate_sql(prompt: str):
         return_tensors="pt"
     )
 
+    input_length = inputs.shape[-1]   # length of prompt tokens
+
     with torch.inference_mode():
         outputs = model.generate(
             input_ids=inputs,
@@ -61,13 +63,10 @@ def generate_sql(prompt: str):
             pad_token_id=tokenizer.eos_token_id,
         )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # 🔑 Remove the prompt tokens from the output
+    generated_tokens = outputs[0][input_length:]
 
-    # Cleanup
-    if "<|assistant|>" in response:
-        response = response.split("<|assistant|>", 1)[-1].strip()
-    if "<|end|>" in response:
-        response = response.split("<|end|>")[0].strip()
+    response = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
 
     return response