Update app.py

app.py

@@ -2,7 +2,6 @@ import torch
 import gradio as gr
 import multiprocessing
 import os
-import time
 from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
 from peft import get_peft_model, LoraConfig, TaskType
 from datasets import load_dataset
@@ -11,34 +10,56 @@ device = "cpu"
 training_process = None
 log_file = "training_status.log"
 
+# Logging function
 def log_status(message):
     with open(log_file, "w") as f:
         f.write(message)
 
+# Read training status
 def read_status():
     if os.path.exists(log_file):
         with open(log_file, "r") as f:
             return f.read()
     return "⏳ در انتظار شروع ترینینگ..."
 
+# Function to find the text column dynamically
+def find_text_column(dataset):
+    sample = dataset["train"][0]  # Get the first row of the training dataset
+    for column in sample.keys():
+        if isinstance(sample[column], str):  # Find the first text-like column
+            return column
+    return None  # No valid text column found
+
+# Model training function
 def train_model(dataset_url, model_url, epochs):
     try:
         log_status("🚀 در حال بارگیری مدل...")
+        
         tokenizer = AutoTokenizer.from_pretrained(model_url, trust_remote_code=True)
         model = AutoModelForCausalLM.from_pretrained(
             model_url, trust_remote_code=True, torch_dtype=torch.float32, device_map="cpu"
         )
 
         lora_config = LoraConfig(
-            task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=32, lora_dropout=0.1, target_modules=["q_proj", "v_proj"]
+            task_type=TaskType.CAUSAL_LM,
+            r=8,
+            lora_alpha=32,
+            lora_dropout=0.1,
+            target_modules=["q_proj", "v_proj"]
         )
-
         model = get_peft_model(model, lora_config)
         model.to(device)
 
         dataset = load_dataset(dataset_url)
+
+        # Automatically detect the correct text column
+        text_column = find_text_column(dataset)
+        if not text_column:
+            log_status("❌ خطا: ستون متنی در دیتاست یافت نشد!")
+            return
+        
         def tokenize_function(examples):
-            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256)
+            return tokenizer(examples[text_column], truncation=True, padding="max_length", max_length=256)
 
         tokenized_datasets = dataset.map(tokenize_function, batched=True)
         train_dataset = tokenized_datasets["train"]
@@ -61,8 +82,8 @@ def train_model(dataset_url, model_url, epochs):
         )
 
         trainer = Trainer(
-            model=model, 
-            args=training_args, 
+            model=model,
+            args=training_args,
             train_dataset=train_dataset
         )
 
@@ -78,6 +99,7 @@ def train_model(dataset_url, model_url, epochs):
     except Exception as e:
         log_status(f"❌ خطا: {str(e)}")
 
+# Start training in a separate process
 def start_training(dataset_url, model_url, epochs):
     global training_process
     if training_process is None or not training_process.is_alive():
@@ -87,26 +109,24 @@ def start_training(dataset_url, model_url, epochs):
     else:
         return "⚠ ترینینگ در حال اجرا است!"
 
+# Function to update the status
 def update_status():
     return read_status()
 
+# Gradio UI
 with gr.Blocks() as app:
     gr.Markdown("# 🚀 AutoTrain DeepSeek R1 (CPU) - نمایش وضعیت لحظه‌ای")
 
-    dataset_url = gr.Textbox(label="Dataset URL (Hugging Face)", placeholder="مثال: samsum")
-    model_url = gr.Textbox(label="Model URL (Hugging Face)", placeholder="مثال: deepseek-ai/deepseek-r1")
-    epochs = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="تعداد Epochs")
-
-    train_button = gr.Button("شروع ترینینگ")
-    output_text = gr.Textbox(label="وضعیت ترینینگ")
-
-    train_button.click(start_training, inputs=[dataset_url, model_url, epochs], outputs=output_text)
+    with gr.Row():
+        dataset_input = gr.Textbox(label="📂 لینک دیتاست (Hugging Face)")
+        model_input = gr.Textbox(label="🤖 مدل پایه (Hugging Face)")
+        epochs_input = gr.Number(label="🔄 تعداد Epochs", value=3)
 
-    # ✅ نمایش وضعیت لحظه‌ای ترینینگ
-    status_box = gr.Textbox(label="مرحله فعلی ترینینگ", interactive=False)
-    refresh_button = gr.Button("🔄 به‌روزرسانی وضعیت")
+    start_button = gr.Button("🚀 شروع ترینینگ")
+    status_output = gr.Textbox(label="📢 وضعیت ترینینگ", interactive=False)
 
-    refresh_button.click(update_status, inputs=[], outputs=status_box)
+    start_button.click(start_training, inputs=[dataset_input, model_input, epochs_input], outputs=status_output)
+    status_button = gr.Button("🔄 بروزرسانی وضعیت")
+    status_button.click(update_status, outputs=status_output)
 
-app.queue()
-app.launch(server_name="0.0.0.0", server_port=7860, share=True)
+app.launch()