update functions call

app.py

@@ -1,4 +1,3 @@
-
 import os
 from threading import Thread
 from typing import Iterator, List, Tuple, Dict, Any
@@ -6,7 +5,16 @@ from typing import Iterator, List, Tuple, Dict, Any
 import gradio as gr
 import spaces
 import torch
-from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling, TrainerCallback,AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline
+from transformers import (
+    TrainingArguments,
+    Trainer,
+    DataCollatorForLanguageModeling,
+    TrainerCallback,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+    pipeline
+)
 from bs4 import BeautifulSoup
 import requests
 import json
@@ -184,10 +192,14 @@ def process_query(query: str) -> Dict[str, Any]:
     general_query_keywords = ["giải thích", "mô tả", "nói cho tôi biết về", "cái gì là", "cách nào"]
     summarize_keywords = ["tóm tắt", "tóm lại", "khái quát", "ngắn gọn"]
     sentiment_keywords = ["cảm xúc", "tâm trạng", "tâm lý", "phân tích cảm xúc"]
+    train_keywords = ["huấn luyện", "train", "fine-tune", "tinh chỉnh"]
 
     query_lower = query.lower()  # Chuyển truy vấn thành chữ thường để so sánh
-    
-    if any(keyword in query_lower for keyword in web_search_keywords):
+
+    if any(keyword in query_lower for keyword in train_keywords):
+        function_name = "train_model"
+        arguments = {"prompt": query}
+    elif any(keyword in query_lower for keyword in web_search_keywords):
         function_name = "web_search"
         arguments = {"query": query}
     elif any(keyword in query_lower for keyword in summarize_keywords):
@@ -202,7 +214,7 @@ def process_query(query: str) -> Dict[str, Any]:
     else:
         function_name = "hard_query"
         arguments = {"prompt": query}
-    
+
     return {
         "name": function_name,
         "arguments": arguments
@@ -214,7 +226,7 @@ def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: L
     """
     function_name = function_call["name"]
     arguments = function_call["arguments"]
-    
+
     if function_name == "web_search":
         query = arguments["query"]
         yield "🔍 Đang thực hiện tìm kiếm trên web..."
@@ -226,10 +238,10 @@ def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: L
         web_summary = '\n\n'.join([f"🔗 **Liên kết**: {res['link']}\n📝 **Mô tả**: {res['text']}" for res in web_results if res["text"] != "Không thể lấy nội dung."])
         if not web_summary:
             web_summary = "⚠️ Không thể lấy nội dung từ kết quả tìm kiếm."
-        
+
         # Trả về kết quả tìm kiếm cho người dùng
         yield "📄 **Kết quả tìm kiếm:**\n" + web_summary
-    
+
     elif function_name == "summarize_query":
         # Khi người dùng yêu cầu tóm tắt, hệ thống sẽ thực hiện tìm kiếm và sau đó tóm tắt kết quả
         query = arguments["prompt"]
@@ -247,13 +259,21 @@ def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: L
         yield "📝 Đang tóm tắt thông tin..."
         summary = summarize_text(combined_text)
         yield "📄 **Tóm tắt:**\n" + summary
-    
+
     elif function_name == "sentiment_analysis":
         prompt_text = arguments["prompt"]
         yield "📊 Đang phân tích tâm lý..."
         sentiment = analyze_sentiment(prompt_text)
         yield sentiment
-    
+
+    elif function_name == "train_model":
+        prompt_text = arguments["prompt"]
+        yield "🛠️ Đang bắt đầu quá trình huấn luyện..."
+        # Gọi hàm huấn luyện
+        training_iterator = train_model()
+        for response in training_iterator:
+            yield response
+
     elif function_name in ["general_query", "hard_query"]:
         prompt_text = arguments["prompt"]
         yield "🤖 Đang tạo phản hồi..."
@@ -269,7 +289,7 @@ def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: L
         )
         for response in response_generator:
             yield response
-    
+
     else:
         yield "⚠️ Lời gọi hàm không được nhận dạng."
 
@@ -291,23 +311,23 @@ def generate(
     # 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":
-        continuous_training(total_steps=300, steps_per_call=50)
         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,
@@ -319,7 +339,7 @@ def generate(
         top_k=top_k,
         repetition_penalty=repetition_penalty
     )
-    
+
     for response in response_iterator:
         yield response
 
@@ -334,6 +354,7 @@ EXAMPLES = [
     ["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
@@ -380,10 +401,11 @@ chat_interface = gr.ChatInterface(
     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 và phân tích tâm lý.",
+    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"
@@ -453,15 +475,19 @@ class SaveCheckpointCallback(TrainerCallback):
         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
+            trainer = kwargs.get('trainer')  # Sử dụng get để tránh KeyError
+            if trainer:
+                trainer.save_model(checkpoint_path)
+            else:
+                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=30, queue=False)
-def run_training():
+@spaces.GPU(duration=60, queue=False)  # Tăng duration lên 60 giây
+def run_training(steps_per_call=5):
     """
     Hàm huấn luyện mô hình sử dụng GPU với thời gian hạn chế.
+    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(
@@ -488,15 +514,15 @@ def run_training():
         per_device_train_batch_size=4,
         per_device_eval_batch_size=4,
         gradient_accumulation_steps=8,
-        num_train_epochs=3,
-        max_steps=5,  # Đặt max_steps tại đây
+        num_train_epochs=1,  # Giới hạn epochs để đảm bảo chỉ huấn luyện 5 bước
+        max_steps=steps_per_call,  # Đặ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
+        logging_steps=1,  # 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
+        eval_steps=steps_per_call,  # Đánh giá sau mỗi 5 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_steps=steps_per_call,  # Lưu checkpoint sau mỗi 5 bước
         save_total_limit=5,  # Giới hạn số lượng checkpoint lưu trữ
         fp16=True,
         report_to="none",
@@ -518,7 +544,7 @@ def run_training():
         eval_dataset=final_dataset['validation'],
         tokenizer=tokenizer,
         data_collator=data_collator,
-        callbacks=[SaveCheckpointCallback()],  # Thêm callback
+        callbacks=[SaveCheckpointCallback()],  # Thêm callback đã sửa đổi
     )
     
     # Kiểm tra nếu có checkpoint
@@ -534,30 +560,16 @@ def run_training():
     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 Huấn Luyện
+def train_model():
     """
-    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.
+    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.
     """
-    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
+    # Gọi hàm huấn luyện với steps_per_call=5
+    result = run_training(steps_per_call=5)
+    yield result
+
+# ---------------------------- Giao Diện Gradio ---------------------------- #
 
 # Tạo giao diện chính của Gradio với CSS tùy chỉnh
 with gr.Blocks(css="""
@@ -585,5 +597,7 @@ with gr.Blocks(css="""
     gr.DuplicateButton(value="Nhân bản Không gian để sử dụng riêng tư", elem_id="duplicate-button")  # Nút nhân bản không gian
     chat_interface.render()  # Hiển thị giao diện trò chuyện
 
+# ---------------------------- Khởi Chạy Ứng Dụng ---------------------------- #
+
 if __name__ == "__main__":
-    demo.queue(max_size=30).launch()  # Khởi chạy ứng dụng Gradio với hàng đợi kích thước tối đa là 20
+    demo.queue(max_size=30).launch()  # Khởi chạy ứng dụng Gradio với hàng đợi kích thước tối đa là 30