Update app.py

app.py

@@ -43,8 +43,8 @@ training_args = TrainingArguments(
     output_dir="./results",
     eval_strategy="epoch",
     learning_rate=2e-5,
-    per_device_train_batch_size=2,  # Reduced batch size
-    per_device_eval_batch_size=2,   # Reduced batch size
+    per_device_train_batch_size=4,  # Reduced batch size
+    per_device_eval_batch_size=4,   # Reduced batch size
     num_train_epochs=1,
     weight_decay=0.01,
     report_to="none",  # Disables wandb logging
@@ -63,8 +63,8 @@ trainer = Trainer(
 trainer.train()
 
 # Save the model
-model.save_pretrained("./fine-tuned-llama")
-tokenizer.save_pretrained("./fine-tuned-llama")
+model.save_pretrained("./fine-tuned-gpt2")
+tokenizer.save_pretrained("./fine-tuned-gpt2")
 
 # Evaluate the model
 #results = trainer.evaluate()
@@ -73,7 +73,7 @@ tokenizer.save_pretrained("./fine-tuned-llama")
 # Create a Gradio interface for text generation
 def generate_text(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate([**inputs)
+    outputs = model.generate(inputs["input_ids"], max_length=50, num_return_sequences=1)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 iface = gr.Interface(fn=generate_text, inputs="text", outputs="text")