Update app.py

app.py

@@ -1,10 +1,11 @@
+import gradio as gr
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 
 # Load the dataset
 dataset = load_dataset("json", data_files="dataset.jsonl")
 
-# Load the model and tokenizer
+# Load the pre-trained model and tokenizer
 model_name = "Salesforce/codegen-2B-multi"
 model = AutoModelForCausalLM.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -21,13 +22,14 @@ training_args = TrainingArguments(
     overwrite_output_dir=True,
     evaluation_strategy="epoch",
     learning_rate=5e-5,
-    per_device_train_batch_size=4,
+    per_device_train_batch_size=2,
     num_train_epochs=3,
     save_strategy="epoch",
     logging_dir="./logs",
+    logging_strategy="epoch",
 )
 
-# Train the model
+# Trainer setup
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -35,7 +37,29 @@ trainer = Trainer(
     eval_dataset=tokenized_dataset["train"],
 )
 
+# Train the model
 trainer.train()
-trainer.save_model("./fine_tuned_codegen")
-tokenizer.save_pretrained("./fine_tuned_codegen")
-print("Training complete. Model saved.")
+
+# Save the fine-tuned model
+trainer.save_model("./fine_tuned_model")
+tokenizer.save_pretrained("./fine_tuned_model")
+
+# Load the fine-tuned model for inference
+fine_tuned_model = AutoModelForCausalLM.from_pretrained("./fine_tuned_model")
+fine_tuned_tokenizer = AutoTokenizer.from_pretrained("./fine_tuned_model")
+
+# Define a Gradio interface for testing the model
+def generate_cypress_code(prompt):
+    inputs = fine_tuned_tokenizer(prompt, return_tensors="pt")
+    outputs = fine_tuned_model.generate(inputs["input_ids"], max_length=150, num_return_sequences=1)
+    return fine_tuned_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Launch the Gradio interface
+interface = gr.Interface(
+    fn=generate_cypress_code,
+    inputs="text",
+    outputs="text",
+    title="Cypress Test Generator",
+    description="Enter a description of the test you want to generate Cypress code for.",
+)
+interface.launch()