Create app.py

app.py

@@ -0,0 +1,96 @@
+import json
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from torch.utils.data import Dataset
+import os
+
+# Step 1: Load and Preprocess Data
+class SpiderDataset(Dataset):
+    def __init__(self, file_paths, tokenizer, max_length=128):
+        self.data = []
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+        for file_path in file_paths:
+            with open(file_path, 'r') as f:
+                self.data.extend(json.load(f))
+        
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        question = item['question']
+        sql_query = item['query']
+
+        # Tokenize inputs and labels
+        input_encoding = self.tokenizer(
+            question, 
+            max_length=self.max_length, 
+            padding="max_length", 
+            truncation=True, 
+            return_tensors="pt"
+        )
+
+        output_encoding = self.tokenizer(
+            sql_query, 
+            max_length=self.max_length, 
+            padding="max_length", 
+            truncation=True, 
+            return_tensors="pt"
+        )
+
+        # Prepare inputs and labels
+        input_ids = input_encoding['input_ids'].squeeze()
+        labels = output_encoding['input_ids'].squeeze()
+
+        return {
+            "input_ids": input_ids,
+            "labels": labels
+        }
+
+# Step 2: Initialize Tokenizer and Model
+tokenizer = GPT2Tokenizer.from_pretrained("distilgpt2")
+tokenizer.pad_token = tokenizer.eos_token  # Set pad token
+
+# Load model with language model head
+model = GPT2LMHeadModel.from_pretrained("distilgpt2")
+
+# Step 3: Load Datasets
+# Assuming the files are in a directory called `space/dataset`
+file_paths = [
+    "space/dataset/train_others.json",
+    "space/dataset/dev.json",
+    "space/dataset/train_spider.json",
+    "space/dataset/test.json"
+]
+train_dataset = SpiderDataset(file_paths, tokenizer)
+
+# Step 4: Define Training Arguments
+training_args = TrainingArguments(
+    output_dir="./distilgpt2-sql-converter",
+    evaluation_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=4,
+    per_device_eval_batch_size=4,
+    num_train_epochs=3,
+    weight_decay=0.01,
+    logging_dir="./logs",
+    save_total_limit=2,
+)
+
+# Step 5: Initialize Trainer with Data Collator
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    data_collator=data_collator,
+)
+
+# Step 6: Train the Model
+trainer.train()
+
+# Step 7: Save the Model and Tokenizer
+model.save_pretrained("./distilgpt2-sql-converter")
+tokenizer.save_pretrained("./distilgpt2-sql-converter")