update app.py

app.py

@@ -49,10 +49,8 @@ custom_role_conversions=None,
 
 # Import tool from Hub
 image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
-    
 agent = CodeAgent(
     model=model,
     tools=[final_answer], ## add your tools here (don't remove final answer)
@@ -64,6 +62,44 @@ agent = CodeAgent(
     description=None,
     prompt_templates=prompt_templates
 )
+GradioUI(agent).launch() 
+
+
+training_args = TrainingArguments(
+    output_dir="./lora-finetuned",
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
+    save_steps=500,
+    logging_dir="./logs",
+    num_train_epochs=3,
+    save_total_limit=2,
+    fp16=True
+)
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets["train"],
+)
+trainer.train() 
+
 
 
-GradioUI(agent).launch()
+
+training_args = TrainingArguments(
+    output_dir="./lora-finetuned",
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
+    save_steps=500,
+    logging_dir="./logs",
+    num_train_epochs=3,
+    save_total_limit=2,
+    fp16=True
+)
+
+trainer=Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets["train"],
+)
+trainer.train()
+