Update app.py

app.py

@@ -26,10 +26,74 @@ def process(action, base_model_id, dataset, system_prompt, user_prompt, schema):
         result = prompt_model(fine_tuned_model_id, system_prompt, user_prompt, schema)
     return result
 
+# Preprocess the dataset
+def preprocess(examples):
+    model_inputs = tokenizer(examples["text"], text_target=examples["sql"], max_length=512, truncation=True)
+    return model_inputs
+
 def fine_tune_model(base_model_id, dataset):
-    tokenizer = download_model(base_model_id)
-    fine_tuned_model_id = upload_model(base_model_id, tokenizer)
-    return fine_tuned_model_id
+#    tokenizer = download_model(base_model_id)
+#    fine_tuned_model_id = upload_model(base_model_id, tokenizer)
+#    return fine_tuned_model_id
+    # Load the dataset
+    dataset = load_dataset("gretelai/synthetic_text_to_sql")
+    
+    # Load pre-trained model and tokenizer
+    model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+    dataset = dataset.map(preprocess, batched=True)
+    
+    # Split dataset to training and validation sets
+    train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))  # Adjust the range as needed
+    val_dataset = dataset["test"].shuffle(seed=42).select(range(100))  # Adjust the range as needed
+    
+    # Set training arguments
+    training_args = Seq2SeqTrainingArguments(
+        output_dir="./results",
+        num_train_epochs=3,  # Adjust as needed
+        per_device_train_batch_size=16,
+        per_device_eval_batch_size=64,
+        warmup_steps=500,
+        weight_decay=0.01,
+        logging_dir="./logs",
+        save_total_limit=2,
+        save_steps=500,
+        eval_steps=500,
+        metric_for_best_model="accuracy",
+        greater_is_better=True,
+        save_on_each_node=True,
+        load_best_model_at_end=True,
+    )
+    
+    # Create Trainer instance
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        compute_metrics=lambda pred: {"accuracy": torch.sum(pred.label_ids == pred.predictions.argmax(-1))},
+    )
+    
+    # Train the model
+    trainer.train()
+
+    # Save the trained model
+    trainer.save_model("./fine_tuned_model")
+
+    # Create a repository object
+    repo = Repository(
+        local_dir="./fine_tuned_model",
+        repo_type="model",
+        repo_id="bstraehle/Meta-Llama-3.1-8B-Instruct-text-to-sql",
+    )
+    
+    # Login to the Hugging Face hub
+    repo.login(token=os.environ["HF_TOKEN"])
+    
+    # Push the model to the hub
+    repo.push_to_hub(commit_message="Initial commit")
 
 def prompt_model(model_id, system_prompt, user_prompt, schema):
     pipe = pipeline("text-generation",