Update app.py

app.py

@@ -1,5 +1,38 @@
-# Load model directly
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("HridaAI/Hrida-T2SQL-3B-128k-V0.1", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("HridaAI/Hrida-T2SQL-3B-128k-V0.1", trust_remote_code=True)
+# Define the model and tokenizer
+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)
+); 
+
+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)
+
+# Generate the output
+outputs = model.generate(inputs, max_length=300)
+print(tokenizer.decode(outputs[0]))