update new feature datasets

app.py

@@ -1,6 +1,9 @@
+
 import os
 from threading import Thread
 from typing import Iterator, List, Tuple, Dict, Any
+import uuid
+import json
 
 import gradio as gr
 import spaces
@@ -8,8 +11,9 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline
 from bs4 import BeautifulSoup
 import requests
-import json
 from functools import lru_cache
+# from huggingface_hub import HfApi, HfFolder
+from datasets import load_dataset, DatasetDict, Dataset, concatenate_datasets
 
 # ---------------------------- Cấu Hình ---------------------------- #
 
@@ -41,6 +45,60 @@ 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")
 
+# ---------------------------- Thiết lập Bộ nhớ Sử dụng Huggingface Datasets ---------------------------- #
+
+HF_DATASET = os.getenv("HF_DATASET")  # Đảm bảo bạn đã set biến môi trường này "your_username/chat_memory" Thay đổi theo tên của bạn
+
+def initialize_dataset():
+    """
+    Khởi tạo Dataset trên Huggingface Hub nếu chưa tồn tại.
+    """
+    try:
+        dataset = load_dataset(HF_DATASET)
+    except:
+        # Tạo Dataset mới nếu chưa tồn tại
+        dataset = DatasetDict({
+            "conversations": Dataset.from_dict({
+                "user_id": [],
+                "messages": []
+            })
+        })
+        dataset.push_to_hub(HF_DATASET, private=True)
+    return dataset
+
+def save_conversation(user_id: str, messages: List[Tuple[str, str]]):
+    """
+    Lưu cuộc hội thoại của người dùng vào Dataset.
+    """
+    dataset = load_dataset(HF_DATASET)
+    # Chuyển đổi cuộc hội thoại thành định dạng JSON
+    messages_json = json.dumps(messages)
+    new_entry = {
+        "user_id": user_id,
+        "messages": messages_json
+    }
+    # Tạo Dataset từ entry mới
+    new_dataset = Dataset.from_dict(new_entry)
+    # Kết hợp với Dataset hiện tại
+    updated_dataset = concatenate_datasets([dataset["conversations"], new_dataset])
+    # Đẩy lên Hub
+    updated_dataset.push_to_hub(HF_DATASET, private=True)
+
+def load_conversation(user_id: str) -> List[Tuple[str, str]]:
+    """
+    Truy xuất cuộc hội thoại của người dùng từ Dataset.
+    """
+    dataset = load_dataset(HF_DATASET)
+    # Tìm entry theo user_id
+    user_data = dataset["conversations"].filter(lambda x: x["user_id"] == user_id)
+    if len(user_data) == 0:
+        return []
+    messages_json = user_data["messages"][0]
+    return json.loads(messages_json)
+
+# Khởi tạo Dataset
+initialize_dataset()
+
 # ---------------------------- Định Nghĩa Hàm ---------------------------- #
 
 @lru_cache(maxsize=128)
@@ -103,7 +161,7 @@ def summarize_text(text: str, max_length: int = 150) -> str:
     ]
     input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
     input_ids = input_ids.to(device)
-    
+
     summary_streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     summary_kwargs = {
         "input_ids": input_ids,
@@ -115,7 +173,7 @@ def summarize_text(text: str, max_length: int = 150) -> str:
     }
     t = Thread(target=model.generate, kwargs=summary_kwargs)
     t.start()
-    
+
     summary = ""
     for new_text in summary_streamer:
         summary += new_text
@@ -128,26 +186,43 @@ def analyze_sentiment(text: str) -> str:
     score = result[0]['score']
     return f"🟢 **Tâm lý**: {sentiment} (Điểm: {score:.2f})"
 
-def generate_response(prompt: str, chat_history: List[Tuple[str, str]], max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float) -> Iterator[str]:
+def generate_response(prompt: str, chat_history: List[Tuple[str, str]], user_id: str, max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float) -> Iterator[str]:
     """
     Tạo phản hồi sử dụng mô hình Llama cục bộ theo chế độ streaming.
     """
-    # Xây dựng lịch sử cuộc trò chuyện
-    conversation = []
-    for user, assistant in chat_history:
-        conversation.extend([
-            {"role": "user", "content": user},
-            {"role": "assistant", "content": assistant},
+    # Lấy lịch sử từ Dataset
+    conversation = load_conversation(user_id)
+    # Chuyển đổi lịch sử thành định dạng mà mô hình hiểu
+    conversation_formatted = []
+    for user_msg, assistant_msg in conversation:
+        conversation_formatted.extend([
+            {"role": "user", "content": user_msg},
+            {"role": "assistant", "content": assistant_msg},
         ])
-    conversation.append({"role": "user", "content": prompt})  # Thêm tin nhắn của người dùng
+    conversation_formatted.append({"role": "user", "content": prompt})  # Thêm tin nhắn của ngư��i dùng
+
+    # Kiểm tra độ dài và sử dụng bản tóm tắt nếu cần
+    if len(conversation_formatted) > 50:  # Giới hạn số lượng tin nhắn, điều chỉnh tùy nhu cầu
+        summary = summarize_text(" ".join([msg["content"] for msg in conversation_formatted]))
+        # Lưu bản tóm tắt vào Dataset
+        new_messages = [("system", summary)]
+        save_conversation(user_id, new_messages)
+        # Giữ lại phần mới nhất
+        conversation_formatted = [{"role": "system", "content": summary}] + conversation_formatted[-25:]
     
     # Chuẩn bị input_ids từ tokenizer
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
+    input_ids = tokenizer.apply_chat_template(conversation_formatted, add_generation_prompt=True, return_tensors="pt")
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]  # Cắt input nếu quá dài
-        gr.Warning(f"Đã cắt bỏ phần cuộc trò chuyện vì vượt quá {MAX_INPUT_TOKEN_LENGTH} token.")
+        # Sử dụng bản tóm tắt từ bộ nhớ
+        summary = summarize_text(" ".join([msg["content"] for msg in conversation_formatted]))
+        conversation_formatted = [{"role": "system", "content": summary}] + conversation_formatted[-(MAX_INPUT_TOKEN_LENGTH // 2):]
+        input_ids = tokenizer.apply_chat_template(conversation_formatted, add_generation_prompt=True, return_tensors="pt")
+        # Lưu lại bản tóm tắt
+        new_messages = [("system", summary)]
+        save_conversation(user_id, new_messages)
+
     input_ids = input_ids.to(device)  # Di chuyển input tới thiết bị
-    
+
     # Khởi tạo streamer để nhận văn bản được tạo ra theo thời gian thực
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = {
@@ -163,13 +238,17 @@ def generate_response(prompt: str, chat_history: List[Tuple[str, str]], max_new_
     }
     t = Thread(target=model.generate, kwargs=generate_kwargs)  # Tạo luồng để sinh văn bản
     t.start()
-    
+
     # Stream văn bản được tạo ra
     outputs = []
     for text in streamer:
         outputs.append(text)
         yield "".join(outputs)
 
+    # Lưu phản hồi vào Dataset
+    response = "".join(outputs)
+    save_conversation(user_id, [(prompt, response)])
+
 @lru_cache(maxsize=128)
 def process_query(query: str) -> Dict[str, Any]:
     """
@@ -180,9 +259,10 @@ 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"]
+    topic_keywords = ["chủ đề", "bàn về", "về"]
 
     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):
         function_name = "web_search"
         arguments = {"query": query}
@@ -192,25 +272,28 @@ def process_query(query: str) -> Dict[str, Any]:
     elif any(keyword in query_lower for keyword in sentiment_keywords):
         function_name = "sentiment_analysis"
         arguments = {"prompt": query}
+    elif any(keyword in query_lower for keyword in topic_keywords):
+        function_name = "new_topic"
+        arguments = {"topic": query}
     elif any(keyword in query_lower for keyword in general_query_keywords):
         function_name = "general_query"
         arguments = {"prompt": query}
     else:
         function_name = "hard_query"
         arguments = {"prompt": query}
-    
+
     return {
         "name": function_name,
         "arguments": arguments
     }
 
-def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: List[Tuple[str, str]], max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float) -> Iterator[str]:
+def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: List[Tuple[str, str]], user_id: str, max_new_tokens: int, temperature: float, top_p: float, top_k: int, repetition_penalty: float) -> Iterator[str]:
     """
     Thực thi hàm phù hợp dựa trên lời gọi hàm.
     """
     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..."
@@ -222,10 +305,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"]
@@ -242,14 +325,22 @@ def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: L
         # Tóm tắt nội dung đã lấy
         yield "📝 Đang tóm tắt thông tin..."
         summary = summarize_text(combined_text)
+        # Lưu tóm tắt vào Dataset
+        save_conversation(user_id, [("tóm tắt", summary)])
         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 == "new_topic":
+        topic = arguments["topic"]
+        # Lưu chủ đề mới vào Dataset
+        save_conversation(user_id, [("chủ đề", f"Chủ đề mới: {topic}")])
+        yield f"🆕 Đã chuyển sang chủ đề mới: {topic}"
+
     elif function_name in ["general_query", "hard_query"]:
         prompt_text = arguments["prompt"]
         yield "🤖 Đang tạo phản hồi..."
@@ -257,6 +348,7 @@ def handle_functions(function_call: Dict[str, Any], prompt: str, chat_history: L
         response_generator = generate_response(
             prompt=prompt_text,
             chat_history=chat_history,
+            user_id=user_id,
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             top_p=top_p,
@@ -265,12 +357,23 @@ 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."
 
 # ---------------------------- Giao Diện Gradio ---------------------------- #
 
+def get_user_id():
+    """
+    Tạo hoặc lấy user_id từ session state của Gradio.
+    Sử dụng cookie hoặc thông tin định danh tạm thời.
+    """
+    # Gradio hiện không hỗ trợ session state natively, cần sử dụng workaround
+    # Dưới đây là cách tạo user_id tạm thời cho mỗi phiên
+    if "user_id" not in gr.get_session_state():
+        gr.get_session_state()["user_id"] = str(uuid.uuid4())
+    return gr.get_session_state()["user_id"]
+
 @spaces.GPU(duration=15, queue=False)
 def generate(
     message: str,
@@ -286,10 +389,13 @@ 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..."
-    
+
+    # Lấy user_id từ session
+    user_id = get_user_id()
+
     # 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."
@@ -297,23 +403,27 @@ def generate(
         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"] == "new_topic":
+        yield "🛠️ Đã chọn chức năng: Chủ đề mới."
+    elif function_call["name"] in ["general_query", "hard
     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,
+        user_id=user_id,  # Sử dụng user_id để quản lý dữ liệu theo người dùng
         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