From dee7e6e9601b5f09081373613617d026e2f65e15 Mon Sep 17 00:00:00 2001
From: Tony Tran <thanhtan.tran@gmail.com>
Date: Mon, 20 Apr 2026 20:39:43 +0700
Subject: [PATCH] =?UTF-8?q?b=E1=BB=95=20sung=20t=E1=BA=A3i=20model=20lu?=
 =?UTF-8?q?=C3=B4n?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server.py | 64 +++++++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 57 insertions(+), 7 deletions(-)

diff --git a/server.py b/server.py
index 3a5025b..d433747 100644
--- a/server.py
+++ b/server.py
@@ -22,18 +22,45 @@ TEMPLATE_DIR = Path(__file__).parent / "templates"
 AVAILABLE_MODELS = {
     "gemma-4-E2B-it": {
         "file": "gemma-4-E2B-it.litertlm",
-        "repo": "google/gemma-4-E2B-it",
+        "repo": "litert-community/gemma-4-E2B-it-litert-lm",
         "desc": "Gemma 4 Edge 2B — nhỏ hơn, nhanh hơn",
     },
     "gemma-4-E4B-it": {
         "file": "gemma-4-E4B-it.litertlm",
-        "repo": "google/gemma-4-E4B-it",
+        "repo": "litert-community/gemma-4-E4B-it-litert-lm",
         "desc": "Gemma 4 Edge 4B — thông minh hơn, chậm hơn",
     },
 }
 
 # ── CLI: chọn model khi khởi động ────────────────────────────────────────────
 
+def download_model(repo: str, local_dir: Path) -> bool:
+    """Tải model từ Hugging Face về local."""
+    try:
+        import subprocess
+        print(f"\n  Đang tải model từ {repo}...")
+        print(f"  Vui lòng đợi, quá trình này có thể mất vài phút...\n")
+        
+        cmd = [
+            "huggingface-cli", "download", repo,
+            "--include", "*.litertlm",
+            "--local-dir", str(local_dir)
+        ]
+        
+        result = subprocess.run(cmd, check=True, capture_output=False, text=True)
+        print(f"\n  ✓ Tải model thành công!")
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"\n  ✗ Lỗi khi tải model: {e}")
+        return False
+    except FileNotFoundError:
+        print(f"\n  ✗ Không tìm thấy huggingface-cli.")
+        print(f"  Cài đặt bằng lệnh: pip install huggingface-hub")
+        return False
+    except Exception as e:
+        print(f"\n  ✗ Lỗi không xác định: {e}")
+        return False
+
 def select_model() -> Path:
     print("\n" + "="*52)
     print("  LiteRT-LM Server — Chọn model")
@@ -58,15 +85,38 @@ def select_model() -> Path:
 
                 if not model_path.exists():
                     print(f"\n  Model chưa có trong thư mục models/")
-                    print(f"  Tải về bằng lệnh:\n")
+                    print(f"  Lệnh tải thủ công:\n")
                     print(f"    huggingface-cli download {info['repo']} \\")
                     print(f"      --include '*.litertlm' \\")
                     print(f"      --local-dir models/\n")
-                    retry = input("  Chọn model khác? (y/n): ").strip().lower()
-                    if retry == "y":
-                        continue
+                    
+                    download_choice = input("  Bạn muốn tải model ngay bây giờ? (y/n): ").strip().lower()
+                    if download_choice == "y":
+                        if download_model(info['repo'], MODELS_DIR):
+                            # Kiểm tra lại xem file đã tồn tại chưa
+                            if model_path.exists():
+                                print(f"\n  Đã chọn: {key}")
+                                print(f"  Path: {model_path}\n")
+                                return model_path
+                            else:
+                                print(f"\n  ✗ Không tìm thấy file model sau khi tải.")
+                                retry = input("  Chọn model khác? (y/n): ").strip().lower()
+                                if retry == "y":
+                                    continue
+                                else:
+                                    sys.exit(0)
+                        else:
+                            retry = input("\n  Chọn model khác? (y/n): ").strip().lower()
+                            if retry == "y":
+                                continue
+                            else:
+                                sys.exit(0)
                     else:
-                        sys.exit(0)
+                        retry = input("  Chọn model khác? (y/n): ").strip().lower()
+                        if retry == "y":
+                            continue
+                        else:
+                            sys.exit(0)
 
                 print(f"\n  Đã chọn: {key}")
                 print(f"  Path: {model_path}\n")