join code to app

app.py

@@ -6,12 +6,14 @@ from typing import Iterator, List, Tuple, Dict, Any
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline
+from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling, TrainerCallback,AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline
 from bs4 import BeautifulSoup
 import requests
 import json
 from functools import lru_cache
-from checkpoint import continuous_training
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
+import time
 
 # ---------------------------- Cấu Hình ---------------------------- #
 
@@ -382,6 +384,181 @@ chat_interface = gr.ChatInterface(
     theme="default",  # Có thể thay đổi theme để giao diện đẹp hơn
 )
 
+
+# Đường dẫn lưu checkpoint
+CHECKPOINT_DIR = "./checkpoints"
+if not os.path.exists(CHECKPOINT_DIR):
+    os.makedirs(CHECKPOINT_DIR)
+
+# Tải Dataset (CPU)
+dataset = load_dataset('vntc/wiki-mini-corpus')
+
+# Chia Dataset thành train và validation (CPU)
+split_dataset = dataset['train'].train_test_split(test_size=0.1, seed=42)
+train_dataset = split_dataset['train']
+validation_dataset = split_dataset['test']
+
+# Tiền Xử Lý Văn Bản (CPU)
+def preprocess_function(examples):
+    passages = [passage.lower().strip() for passage in examples['passage']]
+    return {'passage': passages}
+
+processed_train = train_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])
+processed_validation = validation_dataset.map(preprocess_function, batched=True, remove_columns=['id', 'metadata'])
+
+# Tokenization (CPU)
+model_name = "meta-llama/Llama-3.2-3B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Đảm bảo tokenizer có pad_token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+def tokenize_function(examples):
+    return tokenizer(
+        examples['passage'],
+        padding='max_length',
+        truncation=True,
+        max_length=512,
+    )
+
+tokenized_train = processed_train.map(tokenize_function, batched=True)
+tokenized_validation = processed_validation.map(tokenize_function, batched=True)
+
+# Thêm trường 'labels' (CPU)
+def add_labels(examples):
+    examples['labels'] = examples['input_ids'].copy()
+    return examples
+
+tokenized_train = tokenized_train.map(add_labels, batched=True)
+tokenized_validation = tokenized_validation.map(add_labels, batched=True)
+
+# Loại bỏ các cột không cần thiết (CPU)
+tokenized_train = tokenized_train.remove_columns(['passage'])
+tokenized_validation = tokenized_validation.remove_columns(['passage'])
+
+# Định dạng dữ liệu cho PyTorch (CPU)
+tokenized_train.set_format('torch')
+tokenized_validation.set_format('torch')
+
+# Tạo DatasetDict (CPU)
+final_dataset = {
+    'train': tokenized_train,
+    'validation': tokenized_validation
+}
+
+# Định Nghĩa TrainerCallback để Lưu Checkpoint Nhanh Hơn
+class SaveCheckpointCallback(TrainerCallback):
+    def on_step_end(self, args, state, control, **kwargs):
+        if state.global_step % args.save_steps == 0 and state.global_step != 0:
+            checkpoint_path = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
+            print(f"Lưu checkpoint tại: {checkpoint_path}")
+            trainer = kwargs['trainer']  # Truy cập trainer từ kwargs
+            trainer.save_model(checkpoint_path)
+        return control  # Trả về đối tượng control hiện tại
+
+# Định Nghĩa Hàm Huấn Luyện với Decorator @spaces.GPU
+@spaces.GPU(duration=15, queue=False)
+def run_training():
+    """
+    Hàm huấn luyện mô hình sử dụng GPU với thời gian hạn chế.
+    """
+    # Tải và Cấu Hình Mô Hình với LoRA (GPU)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        load_in_8bit=False
+    )
+    
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=32,
+        target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
+        lora_dropout=0.1,
+        bias="none",
+    )
+    
+    model = get_peft_model(model, lora_config)
+    print(model)
+    
+    # Cấu Hình TrainingArguments (GPU)
+    training_args = TrainingArguments(
+        output_dir=CHECKPOINT_DIR,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        gradient_accumulation_steps=8,
+        num_train_epochs=3,
+        max_steps=50,  # Đặt max_steps tại đây
+        learning_rate=3e-4,
+        weight_decay=0.01,
+        logging_steps=5,  # Giảm số bước logging để theo dõi thường xuyên hơn
+        eval_strategy="steps",  # Đánh giá sau mỗi vài bước
+        eval_steps=5,  # Đánh giá sau mỗi 50 bước
+        save_strategy="steps",  # Lưu checkpoint sau mỗi vài bước
+        save_steps=5,  # Lưu checkpoint sau mỗi 50 bước
+        save_total_limit=5,  # Giới hạn số lượng checkpoint lưu trữ
+        fp16=True,
+        report_to="none",
+        load_best_model_at_end=True,
+    )
+    
+    # Data Collator (GPU)
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer, 
+        mlm=False,  # Vì bạn đang thực hiện Causal LM
+        pad_to_multiple_of=8
+    )
+    
+    # Tạo Trainer (GPU)
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=final_dataset['train'],
+        eval_dataset=final_dataset['validation'],
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        callbacks=[SaveCheckpointCallback()],  # Thêm callback
+    )
+    
+    # Kiểm tra nếu có checkpoint
+    checkpoints = [os.path.join(CHECKPOINT_DIR, d) for d in os.listdir(CHECKPOINT_DIR) if d.startswith('checkpoint')]
+    if checkpoints:
+        latest_checkpoint = max(checkpoints, key=os.path.getctime)
+        print(f"Đang tiếp tục huấn luyện từ checkpoint: {latest_checkpoint}")
+        trainer.train(resume_from_checkpoint=latest_checkpoint)
+    else:
+        trainer.train()
+    
+    # Lưu checkpoint sau khi huấn luyện
+    trainer.save_model(CHECKPOINT_DIR)
+    return "Huấn luyện hoàn tất hoặc đã tiếp tục từ checkpoint."
+
+# Hàm Tự Động Hóa Việc Gọi Lặp Lại Hàm Huấn Luyện
+def continuous_training(total_steps=300, steps_per_call=5):
+    """
+    Hàm tự động gọi lại `run_training` để hoàn thành quá trình huấn luyện.
+    
+    Args:
+        total_steps (int): Tổng số bước huấn luyện mong muốn.
+        steps_per_call (int): Số bước huấn luyện mỗi lần gọi hàm.
+    """
+    steps_done = 0
+    while steps_done < total_steps:
+        print(f"Bắt đầu huấn luyện cho {steps_per_call} bước.")
+        result = run_training()
+        print(result)
+        steps_done += steps_per_call
+        print(f"Đã huấn luyện {steps_done} / {total_steps} bước.")
+        
+        # Kiểm tra nếu đã đạt số bước mong muốn
+        if steps_done >= total_steps:
+            print("Đã hoàn thành toàn bộ quá trình huấn luyện.")
+            break
+        
+        # Chờ một khoảng thời gian trước khi gọi lại (tùy thuộc vào yêu cầu của hệ thống)
+        time.sleep(2)  # Thời gian chờ có thể điều chỉnh
+
 # Tạo giao diện chính của Gradio với CSS tùy chỉnh
 with gr.Blocks(css="""
     .gradio-container {