Update app.py

app.py

@@ -1,38 +1,33 @@
-import gradio
+import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Define the model and tokenizer
+# Load the pre-trained model and tokenizer from Hugging Face Model Hub
 model_id = "HridaAI/Hrida-T2SQL-3B-128k-V0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, trust_remote_code=True)
 
-# Define the context and prompt
-prompt = """
-Answer to the query will be in the form of an SQL query.
-### Context: CREATE TABLE Employees (
-    EmployeeID INT PRIMARY KEY,
-    FirstName VARCHAR(50),
-    LastName VARCHAR(50),
-    Age INT,
-    DepartmentID INT,
-    Salary DECIMAL(10, 2),
-    DateHired DATE,
-    Active BOOLEAN,
-    FOREIGN KEY (DepartmentID) REFERENCES Departments(DepartmentID)
-); 
+# Define the function to generate SQL query from natural language input
+def generate_sql(query):
+    # Tokenize input
+    inputs = tokenizer(query, return_tensors="pt")
+    
+    # Generate the SQL query
+    outputs = model.generate(**inputs, max_new_tokens=256)
+    
+    # Decode the generated output into a string
+    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return sql_query
 
-CREATE TABLE Departments (
-    DepartmentID INT PRIMARY KEY,
-    DepartmentName VARCHAR(100),
-    Location VARCHAR(100)
-); 
-### Input: Write a SQL query to select all the employees who are active.
-### Response:
-"""
-# Prepare the input
-messages = [{"role": "user", "content": prompt}]
-inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=generate_sql,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your natural language question here..."),
+    outputs="text",
+    title="Text to SQL Converter",
+    description="Convert natural language questions into SQL queries using the Hrida-T2SQL-3B model."
+)
 
-# Generate the output
-outputs = model.generate(inputs, max_length=300)
-print(tokenizer.decode(outputs[0]))
+# Launch the interface
+iface.launch()