Update app.py

app.py

@@ -1,4 +1,5 @@
 import torch
+import gradio as gr
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -7,82 +8,98 @@ from transformers import (
     DataCollatorForLanguageModeling
 )
 from datasets import load_dataset
-import os
+import logging
+import sys
+
+# Configure logging
+logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 
 def train():
-    # Load model and tokenizer
-    model_name = "microsoft/phi-2"
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", trust_remote_code=True)
+    try:
+        # Load model and tokenizer
+        model_name = "microsoft/phi-2"
+        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", trust_remote_code=True)
 
-    # Add padding token if missing
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
+        # Add padding token
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
 
-    # Load dataset (update paths as needed)
-    dataset = load_dataset(
-        "csv",
-        data_files={
-            "train": "eswardivi/medical_qa",
-            "validation": "eswardivi/medical_qa"
-        }
-    )
+        # Load dataset
+        dataset = load_dataset(
+            "csv",
+            data_files={
+                "train": "data/train/data.csv",
+                "validation": "data/validation/data.csv"
+            }
+        )
+
+        # Tokenization function
+        def tokenize_function(examples):
+            return tokenizer(
+                examples["text"],
+                padding="max_length",
+                truncation=True,
+                max_length=256,
+                return_tensors="pt",
+            )
 
-    # Tokenization function
-    def tokenize_function(examples):
-        return tokenizer(
-            examples["text"],
-            padding="max_length",
-            truncation=True,
-            max_length=256,
-            return_tensors="pt",
+        tokenized_dataset = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["text"]
         )
 
-    # Preprocess dataset
-    tokenized_dataset = dataset.map(
-        tokenize_function,
-        batched=True,
-        remove_columns=["text"]
-    )
+        # Data collator
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False
+        )
 
-    # Data collator
-    data_collator = DataCollatorForLanguageModeling(
-        tokenizer=tokenizer,
-        mlm=False
-    )
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir="./phi2-results",
+            per_device_train_batch_size=2,
+            per_device_eval_batch_size=2,
+            num_train_epochs=3,
+            logging_dir="./logs",
+            logging_steps=10,
+            fp16=False,
+        )
 
-    # Training arguments
-    training_args = TrainingArguments(
-        output_dir="./phi2-cpu-results",
-        overwrite_output_dir=True,
-        per_device_train_batch_size=2,
-        per_device_eval_batch_size=2,
-        num_train_epochs=3,
-        logging_dir="./logs",
-        logging_steps=100,
-        evaluation_strategy="epoch",
-        save_strategy="epoch",
-        fp16=False,
-        report_to="none",
-    )
+        # Trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset["train"],
+            eval_dataset=tokenized_dataset["validation"],
+            data_collator=data_collator,
+        )
 
-    # Initialize Trainer
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=tokenized_dataset["train"],
-        eval_dataset=tokenized_dataset["validation"],
-        data_collator=data_collator,
-    )
+        # Start training
+        logging.info("Training started...")
+        trainer.train()
+        trainer.save_model("./phi2-trained-model")
+        logging.info("Training completed!")
+
+        return "✅ Training succeeded! Model saved."
 
-    # Start training
-    print("Starting training...")
-    trainer.train()
+    except Exception as e:
+        logging.error(f"Training failed: {str(e)}")
+        return f"❌ Training failed: {str(e)}"
 
-    # Save model
-    trainer.save_model("./phi2-trained-model")
-    tokenizer.save_pretrained("./phi2-trained-model")
-    print("Training complete! Model saved.")
+# Gradio UI
+with gr.Blocks(title="Phi-2 Training") as demo:
+    gr.Markdown("# 🚀 Train Phi-2 on CPU")
+    
+    with gr.Row():
+        start_btn = gr.Button("Start Training", variant="primary")
+        status_output = gr.Textbox(label="Status", interactive=False)
+    
+    start_btn.click(
+        fn=train,
+        outputs=status_output
+    )
 
 if __name__ == "__main__":
-    train()
+    demo.launch(server_name="0.0.0.0", server_port=7860)