Update app.py

app.py

@@ -1,51 +1,26 @@
-import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration, AdamW
-from torch.utils.data import DataLoader, Dataset
-from tqdm import tqdm
-
-# Define your dataset class
-class SpiderDataset(Dataset):
-    def __init__(self, encodings, labels):
-        self.encodings = encodings
-        self.labels = labels
-
-    def __getitem__(self, idx):
-        return {'input_ids': self.encodings[idx], 'labels': self.labels[idx]}
-
-    def __len__(self):
-        return len(self.encodings)
-
-# Load your preprocessed Spider dataset
-train_encodings = # Your preprocessed input encodings for training
-train_labels = # Your preprocessed labels for training
-
-# Create a PyTorch dataset and dataloader
-train_dataset = SpiderDataset(train_encodings, train_labels)
-train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
-
-# Load the pre-trained T5 model
-model = T5ForConditionalGeneration.from_pretrained('t5-base')
-tokenizer = T5Tokenizer.from_pretrained('t5-base')
-
-# Move the model to the GPU if available
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model.to(device)
-
-# Set up the optimizer
-optimizer = AdamW(model.parameters(), lr=5e-5)
-
-# Fine-tune the model
-model.train()
-for epoch in range(3):  # Number of epochs
-    for batch in tqdm(train_loader):
-        optimizer.zero_grad()
-        input_ids = batch['input_ids'].to(device)
-        labels = batch['labels'].to(device)
-        outputs = model(input_ids=input_ids, labels=labels)
-        loss = outputs.loss
-        loss.backward()
-        optimizer.step()
-
-# Save the fine-tuned model
-model.save_pretrained('your_model_directory')
-tokenizer.save_pretrained('your_model_directory')
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+model = AutoModelForSeq2SeqLM.from_pretrained("mrm8488/t5-base-finetuned-wikiSQL")
+
+def generate_sql(query):
+    input_text = "translate English to SQL: " + query
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True)
+    outputs = model.generate(**inputs, max_length=512)
+    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return sql_query
+
+# Create a Gradio interface
+interface = gr.Interface(
+    fn=generate_sql,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your natural language query here..."),
+    outputs="text",
+    title="NL to SQL with T5",
+    description="This model converts natural language queries into SQL. Enter your query!"
+)
+
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()