limit steps

app.py

@@ -46,6 +46,7 @@ model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
     torch_dtype=torch.bfloat16,  # Sử dụng dtype phù hợp để tiết kiệm bộ nhớ
+    load_in_8bit=False 
 )
 model.to(device)  # Di chuyển mô hình tới thiết bị đã chọn
 model.eval()  # Đặt mô hình ở chế độ đánh giá
@@ -53,6 +54,70 @@ model.eval()  # Đặt mô hình ở chế độ đánh giá
 # Khởi tạo pipeline phân tích tâm lý
 sentiment_pipeline = pipeline("sentiment-analysis", model="nlptown/bert-base-multilingual-uncased-sentiment")
 
+# ---------------------------- Tải và Tiền Xử Lý Dữ Liệu ---------------------------- #
+
+# Đườ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 Hàm ---------------------------- #
 
 @lru_cache(maxsize=128)
@@ -293,183 +358,8 @@ def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: L
     else:
         yield "⚠️ Lời gọi hàm không được nhận dạng."
 
-# ---------------------------- Giao Diện Gradio ---------------------------- #
-
-@spaces.GPU(duration=30, queue=False)
-def generate(
-    message: str,
-    chat_history: List[Tuple[str, str]],
-    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
-) -> Iterator[str]:
-    """
-    Hàm chính để xử lý đầu vào của người dùng và tạo phản hồi.
-    """
-    # Thông báo về việc phân tích đầu vào
-    yield "🔍 Đang phân tích truy vấn của bạn..."
-
-    # Xác định hàm nào sẽ được gọi dựa trên tin nhắn của người dùng
-    function_call = process_query(message)
-
-    # Thông báo về hàm được chọn
-    if function_call["name"] == "web_search":
-        yield "🛠️ Đã chọn chức năng: Tìm kiếm trên web."
-    elif function_call["name"] == "summarize_query":
-        yield "🛠️ Đã chọn chức năng: Tóm tắt văn bản."
-    elif function_call["name"] == "sentiment_analysis":
-        yield "🛠️ Đã chọn chức năng: Phân tích tâm lý."
-    elif function_call["name"] == "train_model":
-        yield "🛠️ Đã chọn chức năng: Huấn luyện mô hình."
-    elif function_call["name"] in ["general_query", "hard_query"]:
-        yield "🛠️ Đã chọn chức năng: Trả lời câu hỏi."
-    else:
-        yield "⚠️ Không thể xác định chức năng phù hợp."
-
-    # Xử lý lời gọi hàm và sinh phản hồi tương ứng
-    response_iterator = handle_functions(
-        function_call=function_call,
-        prompt=message,
-        chat_history=chat_history,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        repetition_penalty=repetition_penalty
-    )
-
-    for response in response_iterator:
-        yield response
-
-# Định nghĩa các ví dụ để hướng dẫn người dùng
-EXAMPLES = [
-    ["Xin chào! Bạn khỏe không?"],
-    ["Bạn có thể giải thích ngắn gọn về ngôn ngữ lập trình Python không?"],
-    ["Giải thích cốt truyện của Cô bé Lọ Lem trong một câu."],
-    ["Một người đàn ông cần bao nhiêu giờ để ăn một chiếc máy bay trực thăng?"],
-    ["Viết một bài báo 100 từ về 'Lợi ích của mã nguồn mở trong nghiên cứu AI'"],
-    ["Tìm và cung cấp cho tôi tin tức mới nhất về năng lượng tái tạo."],
-    ["Tìm thông tin về Rạn san hô Great Barrier Reef."],
-    ["Tóm tắt nội dung về trí tuệ nhân tạo."],
-    ["Phân tích tâm lý của đoạn văn sau: Tôi rất vui khi được gặp bạn hôm nay!"],
-    ["Huấn luyện mô hình với dữ liệu mới để cải thiện khả năng hiểu tiếng Việt."],  # Ví dụ mới thêm
-]
-
-# Cấu hình giao diện trò chuyện của Gradio với giao diện đẹp mắt
-chat_interface = gr.ChatInterface(
-    fn=generate,  # Hàm được gọi khi có tương tác từ người dùng
-    additional_inputs=[
-        gr.Slider(
-            label="Số token mới tối đa",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Nhiệt độ",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
-        ),
-        gr.Slider(
-            label="Top-k",
-            minimum=1,
-            maximum=1000,
-            step=1,
-            value=50,
-        ),
-        gr.Slider(
-            label="Hình phạt sự lặp lại",
-            minimum=1.0,
-            maximum=2.0,
-            step=0.05,
-            value=1.2,
-        ),
-    ],
-    stop_btn=None,  # Không có nút dừng
-    examples=EXAMPLES,  # Các ví dụ được hiển thị cho người dùng
-    cache_examples=False,  # Không lưu bộ nhớ cache cho các ví dụ
-    title="🤖 OpenGPT-4o Chatbot",
-    description="Một trợ lý AI mạnh mẽ sử dụng mô hình Llama-3.2 cục bộ với các chức năng tìm kiếm web, tóm tắt văn bản, phân tích tâm lý và huấn luyện mô hình.",
-    theme="default",  # Có thể thay đổi theme để giao diện đẹp hơn
-)
-
 # ---------------------------- Huấn Luyện Mô Hình ---------------------------- #
 
-# Đườ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:
@@ -482,7 +372,6 @@ class SaveCheckpointCallback(TrainerCallback):
                 print("Không thể truy cập 'trainer' từ kwargs.")
         return control
 
-# Định Nghĩa Hàm Huấn Luyện với Decorator @spaces.GPU
 @spaces.GPU(duration=60, queue=False)  # Tăng duration lên 60 giây
 def run_training(steps_per_call=5):
     """
@@ -490,12 +379,12 @@ def run_training(steps_per_call=5):
     Huấn luyện 5 bước mỗi lần gọi.
     """
     # 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
-    )
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     model_name,
+    #     device_map="auto",
+    #     torch_dtype=torch.float16,
+    #     load_in_8bit=False
+    # )
     
     lora_config = LoraConfig(
         r=8,
@@ -560,7 +449,6 @@ def run_training(steps_per_call=5):
     trainer.save_model(CHECKPOINT_DIR)
     return "Huấn luyện hoàn tất hoặc đã tiếp tục từ checkpoint."
 
-# Hàm Huấn Luyện
 def train_model():
     """
     Hàm huấn luyện mô hình, huấn luyện 5 bước mỗi lần và đảm bảo lưu checkpoint.
@@ -571,6 +459,68 @@ def train_model():
 
 # ---------------------------- Giao Diện Gradio ---------------------------- #
 
+# Định nghĩa các ví dụ để hướng dẫn người dùng
+EXAMPLES = [
+    ["Xin chào! Bạn khỏe không?"],
+    ["Bạn có thể giải thích ngắn gọn về ngôn ngữ lập trình Python không?"],
+    ["Giải thích cốt truyện của Cô bé Lọ Lem trong một câu."],
+    ["Một người đàn ông cần bao nhiêu giờ để ăn một chiếc máy bay trực thăng?"],
+    ["Viết một bài báo 100 từ về 'Lợi ích của mã nguồn mở trong nghiên cứu AI'"],
+    ["Tìm và cung cấp cho tôi tin tức mới nhất về năng lượng tái tạo."],
+    ["Tìm thông tin về Rạn san hô Great Barrier Reef."],
+    ["Tóm tắt nội dung về trí tuệ nhân tạo."],
+    ["Phân tích tâm lý của đoạn văn sau: Tôi rất vui khi được gặp bạn hôm nay!"],
+    ["Huấn luyện mô hình với dữ liệu mới để cải thiện khả năng hiểu tiếng Việt."],  # Ví dụ mới thêm
+]
+
+# Cấu hình giao diện trò chuyện của Gradio với giao diện đẹp mắt
+chat_interface = gr.ChatInterface(
+    fn=generate,  # Hàm được gọi khi có tương tác từ người dùng
+    additional_inputs=[
+        gr.Slider(
+            label="Số token mới tối đa",
+            minimum=1,
+            maximum=MAX_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_MAX_NEW_TOKENS,
+        ),
+        gr.Slider(
+            label="Nhiệt độ",
+            minimum=0.1,
+            maximum=4.0,
+            step=0.1,
+            value=0.6,
+        ),
+        gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        ),
+        gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=1000,
+            step=1,
+            value=50,
+        ),
+        gr.Slider(
+            label="Hình phạt sự lặp lại",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.2,
+        ),
+    ],
+    stop_btn=None,  # Không có nút dừng
+    examples=EXAMPLES,  # Các ví dụ được hiển thị cho người dùng
+    cache_examples=False,  # Không lưu bộ nhớ cache cho các ví dụ
+    title="🤖 OpenGPT-4o Chatbot",
+    description="Một trợ lý AI mạnh mẽ sử dụng mô hình Llama-3.2 cục bộ với các chức năng tìm kiếm web, tóm tắt văn bản, phân tích tâm lý và huấn luyện mô hình.",
+    theme="default",  # Có thể thay đổi theme để giao diện đẹp hơn
+)
+
 # Tạo giao diện chính của Gradio với CSS tùy chỉnh
 with gr.Blocks(css="""
     .gradio-container {