Initial training setup

app.py

@@ -0,0 +1,8 @@
+import gradio as gr
+import os
+
+def start_training():
+    os.system("python train.py")
+    return "Training complete! Check model on Hugging Face."
+
+gr.Interface(fn=start_training, inputs=None, outputs="text").launch()

requirements.txt

@@ -1 +1,4 @@
-tk
+transformers==4.54.0
+datasets==4.0.0
+accelerate==1.9.0
+torch>=2.6.0

train.py

@@ -1,52 +1,60 @@
-import json
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
-import huggingface_hub
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling
 
-# 1. Login to Hugging Face (do once)
-# huggingface_hub.login(token="YOUR_HF_TOKEN")
+# Load dataset from Hugging Face Hub
+dataset = load_dataset("Percy3822/quiz_model")
 
-# 2. Load JSONL dataset (local)
-dataset = load_dataset("json", data_files="sample_dataset.jsonl", split="train")
-
-# 3. Convert dict completion → string
+# Preprocess: combine prompt + completion into single string
 def format_for_training(example):
-    completion_str = json.dumps(example["completion"])
-    return {"text": example["prompt"] + "\n" + completion_str}
+    # Convert dict completion to string if needed
+    if isinstance(example["completion"], dict):
+        example["completion"] = str(example["completion"])
+    return {"text": example["prompt"] + "\n" + example["completion"]}
 
 dataset = dataset.map(format_for_training)
 
-# 4. Load small model for low VRAM (good for testing)
-model_name = "distilgpt2"  # Small, works on Hugging Face free GPU
+# Load tokenizer and model (small model for low VRAM)
+model_name = "distilgpt2"  # Small and fast for testing
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
 
-def tokenize_function(examples):
-    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256)
+# Tokenize
+def tokenize(batch):
+    return tokenizer(batch["text"], padding="max_length", truncation=True, max_length=128)
 
-tokenized_dataset = dataset.map(tokenize_function, batched=True)
+dataset = dataset.map(tokenize, batched=True)
 
-# 5. Load model
+# Load model
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# 6. Training Arguments (small batch size for low memory)
+# Data collator
+data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+# Training args
 training_args = TrainingArguments(
     output_dir="./results",
     overwrite_output_dir=True,
-    evaluation_strategy="no",
+    evaluation_strategy="epoch",
+    learning_rate=5e-5,
     per_device_train_batch_size=2,
-    num_train_epochs=2,
-    save_steps=20,
+    num_train_epochs=1,
+    save_strategy="epoch",
+    logging_dir="./logs",
     logging_steps=5,
-    push_to_hub=True,  # Upload to Hugging Face Hub
-    hub_model_id="Percy3822/quiz_model_test"  # Change to your repo name
+    push_to_hub=True,
+    hub_model_id="Percy3822/quiz_model",
 )
 
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=tokenized_dataset,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["train"],  # Use train for eval in testing
+    tokenizer=tokenizer,
+    data_collator=data_collator,
 )
 
-# 7. Train & push to Hugging Face
 trainer.train()
+
+# Push trained model to Hub
 trainer.push_to_hub()