Update

app.py

@@ -9,6 +9,7 @@ from transformers import (
 )
 import gradio as gr
 import torch
+import sqlparse
 
 model_name = os.getenv("HF_MODEL_NAME", None)
 tok = AutoTokenizer.from_pretrained(model_name)
@@ -74,8 +75,16 @@ def bot(input_message: str, db_info="", temperature=0.1, top_p=0.9, top_k=0, rep
 
     # Split the text by "|", and get the last element in the list which should be the final query
     final_query = partial_text.split("|")[1].strip()
+
+    try:
+        # Attempt to format SQL query using sqlparse
+        formatted_query = sqlparse.format(final_query, reindent=True, keyword_case='upper')
+    except Exception:
+        # If formatting fails, use the original, unformatted query
+        formatted_query = final_query
+
     # Convert SQL to markdown (not required, but just to show how to use the markdown module)
-    final_query_markdown = f"```sql\n{final_query}\n```"
+    final_query_markdown = f"```sql\n{formatted_query}\n```"
     return final_query_markdown
 
 with gr.Blocks(css_theme="light") as demo:

requirements.txt

@@ -7,5 +7,4 @@ bitsandbytes
 scipy
 transformers
 accelerate
-huggingface_hub
-markdown
+sqlparse