update

.gitignore

@@ -1,2 +1,3 @@
 gemma-4-E2B-it.*
 gemma-4-E4B-it.*
+models/

README.md

@@ -8,11 +8,18 @@ Chạy mô hình **Gemma 4** trên thiết bị nhúng (Orange Pi 5, Raspberry P
 
 - Python 3.10+
 - [`litert-lm`](https://github.com/google-ai-edge/litert-lm) đã cài và hoạt động
-- Model file: `gemma-4-E2B-it.litertlm` (hoặc model `.litertlm` khác)
 - Thư viện Python:
 
 ```bash
-pip install fastapi uvicorn pydantic
+pip install -r requirements.txt
+```
+
+`requirements.txt`:
+```
+fastapi
+uvicorn
+pydantic
+huggingface_hub
 ```
 
 ---
@@ -23,28 +30,86 @@ pip install fastapi uvicorn pydantic
 .
 ├── app.py               # REST API đơn giản, single-turn
 ├── server.py            # REST API đầy đủ + Web UI, multi-turn sessions
+├── templates/
+│   └── index.html       # Giao diện Web UI (tách riêng khỏi server.py)
+├── models/              # Thư mục chứa các file model .litertlm
+│   ├── gemma-4-E2B-it.litertlm
+│   └── gemma-4-E4B-it.litertlm
+├── requirements.txt
 └── README.md
 ```
 
 ---
 
+## 🤖 Models hỗ trợ
+
+| Model | Repo HuggingFace | Mô tả |
+|-------|-----------------|-------|
+| `gemma-4-E2B-it` | [google/gemma-4-E2B-it](https://huggingface.co/google/gemma-4-E2B-it) | Edge 2B — nhanh hơn, nhẹ hơn |
+| `gemma-4-E4B-it` | [google/gemma-4-E4B-it](https://huggingface.co/google/gemma-4-E4B-it) | Edge 4B — thông minh hơn, nặng hơn |
+
+> **Lưu ý:** Nên dùng bản `-it` (instruction-tuned) cho chat/hỏi đáp. Bản không có `-it` là base model, chỉ predict token tiếp theo, không phù hợp để hội thoại.
+
+### Tải model về
+
+```bash
+# Gemma 4 E2B (nhỏ hơn, ~nhanh hơn)
+huggingface-cli download google/gemma-4-E2B-it \
+  --include '*.litertlm' \
+  --local-dir models/
+
+# Gemma 4 E4B (lớn hơn, ~thông minh hơn)
+huggingface-cli download google/gemma-4-E4B-it \
+  --include '*.litertlm' \
+  --local-dir models/
+```
+
+---
+
 ## 🚀 Hướng dẫn sử dụng
 
-### Bước 1 — Đặt model vào cùng thư mục
+### Bước 1 — Tải ít nhất một model vào thư mục `models/`
 
-```
+Xem lệnh tải ở trên.
-gemma-4-E2B-it.litertlm   ← model file
+
-app.py
+### Bước 2 — Chạy server
-server.py
+
+```bash
+python server.py
 ```
 
-Nếu model ở chỗ khác, sửa biến `MODEL_PATH` ở đầu mỗi file.
+Server sẽ hiển thị menu **chọn model** trước khi khởi động:
+
+```
+====================================================
+  LiteRT-LM Server — Chọn model
+====================================================
+  [1] gemma-4-E2B-it
+      Gemma 4 Edge 2B — nhỏ hơn, nhanh hơn
+      ✓ có sẵn
+
+  [2] gemma-4-E4B-it
+      Gemma 4 Edge 4B — thông minh hơn, chậm hơn
+      ✗ chưa tải
+
+Chọn model (1/2):
+```
+
+Nếu model chưa được tải, server sẽ hiển thị lệnh `huggingface-cli download` tương ứng và hỏi có muốn chọn model khác không.
+
+### Bước 3 — Mở Web UI
+
+```
+http://<địa-chỉ-ip>:8000
+```
+
+Tên model đang chạy sẽ hiển thị trên header của Web UI.
 
 ---
 
 ## 📄 `app.py` — REST API đơn giản
 
-File cơ bản, phù hợp để tích hợp nhanh hoặc test.
+File cơ bản, single-turn, không có menu chọn model. Phù hợp để tích hợp nhanh hoặc test.
 
 ### Chạy
 
@@ -52,8 +117,6 @@ File cơ bản, phù hợp để tích hợp nhanh hoặc test.
 python app.py
 ```
 
-Server khởi động tại `http://0.0.0.0:8000`
-
 ### Endpoint
 
 #### `POST /generate`
@@ -81,40 +144,47 @@ curl -X POST http://localhost:8000/generate \
 
 ## 🖥️ `server.py` — REST API đầy đủ + Web UI
 
-Phiên bản đầy đủ với hỗ trợ multi-turn conversation, quản lý session và giao diện chat trên trình duyệt.
+Phiên bản đầy đủ với chọn model khi khởi động, multi-turn conversation, quản lý session và giao diện chat trên trình duyệt.
-
-### Chạy
-
-```bash
-python server.py
-```
-
-Server khởi động tại `http://0.0.0.0:8000`
 
 ---
 
 ### 🌐 Web UI
 
-Mở trình duyệt và truy cập:
+Mở trình duyệt và truy cập `http://<địa-chỉ-ip>:8000`
-
-```
-http://<địa-chỉ-ip>:8000
-```
 
 Tính năng:
+- **Chọn model khi khởi động** qua menu CLI — tên model hiển thị trực tiếp trên header
 - Giao diện chat trực quan, hỗ trợ tiếng Việt
 - Tự động tạo session khi mở trang
 - Nhớ ngữ cảnh hội thoại trong cùng một session
 - Nút **New** để bắt đầu cuộc trò chuyện mới
 - Nút **Clear** để xóa lịch sử và tạo session mới
 - `Enter` để gửi, `Shift + Enter` để xuống dòng
-- **Render Markdown**: câu trả lời của model hiển thị đúng định dạng (heading, list, code block, table, bold/italic, v.v.)
+- **Render Markdown**: câu trả lời hiển thị đúng định dạng (heading, list, code block, table, bold/italic, v.v.)
-- **Đo tốc độ**: badge `⚡ X tok/s` hiển thị bên dưới mỗi câu trả lời, kèm tổng số token và thời gian xử lý
+- **Đo tốc độ**: badge `⚡ X tok/s` bên dưới mỗi câu trả lời, kèm số token và thời gian xử lý
 
 ---
 
 ### 🔌 REST API
 
+#### `GET /info`
+Trả về thông tin model đang chạy và số session hiện tại.
+
+```bash
+curl http://localhost:8000/info
+```
+
+**Response:**
+
+```json
+{
+  "model": "gemma-4-E2B-it",
+  "sessions": 2
+}
+```
+
+---
+
 #### `POST /generate`
 Single-turn, không nhớ context. Dùng khi chỉ cần hỏi đáp đơn lẻ.
 
@@ -124,6 +194,17 @@ curl -X POST http://localhost:8000/generate \
   -d '{"prompt": "Thủ đô của Việt Nam là gì?"}'
 ```
 
+**Response:**
+
+```json
+{
+  "response": "Thủ đô của Việt Nam là Hà Nội.",
+  "tokens": 12,
+  "elapsed_s": 1.45,
+  "tokens_per_sec": 8.27
+}
+```
+
 ---
 
 #### `POST /chat/new`
@@ -226,31 +307,37 @@ curl -X DELETE http://localhost:8000/chat/$SESSION
 
 ## ⚙️ Cấu hình
 
-Các tham số có thể chỉnh trong đầu mỗi file:
+Các tham số chỉnh trong đầu `server.py`:
 
 | Biến | Mô tả | Mặc định |
 |------|-------|---------|
-| `MODEL_PATH` | Đường dẫn đến file model | `gemma-4-E2B-it.litertlm` |
+| `MODELS_DIR` | Thư mục chứa model | `./models` |
+| `AVAILABLE_MODELS` | Danh sách model + repo HuggingFace | xem trong file |
 | `backend` | Backend inference | `litert_lm.Backend.CPU` |
 | `host` | Địa chỉ lắng nghe | `0.0.0.0` |
 | `port` | Cổng | `8000` |
 
+Để thêm model mới vào menu, bổ sung vào dict `AVAILABLE_MODELS` trong `server.py`:
+
+```python
+AVAILABLE_MODELS = {
+    "gemma-4-E2B-it": {
+        "file": "gemma-4-E2B-it.litertlm",
+        "repo": "google/gemma-4-E2B-it",
+        "desc": "Gemma 4 Edge 2B — nhỏ hơn, nhanh hơn",
+    },
+    "ten-model-moi": {
+        "file": "ten-model-moi.litertlm",
+        "repo": "org/repo-name",
+        "desc": "Mô tả model",
+    },
+}
+```
+
 Để đổi backend sang GPU (nếu thiết bị hỗ trợ):
 
 ```python
-engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.GPU)
+engine = litert_lm.Engine(str(MODEL_PATH), backend=litert_lm.Backend.GPU)
-```
-
----
-
-## 🧪 Test nhanh
-
-```bash
-# Kiểm tra server đang chạy
-curl http://localhost:8000/generate \
-  -X POST \
-  -H "Content-Type: application/json" \
-  -d '{"prompt": "Hello!"}'
 ```
 
 ---
@@ -260,8 +347,9 @@ curl http://localhost:8000/generate \
 - Mỗi session giữ toàn bộ lịch sử hội thoại trong RAM. Nên xóa session khi không dùng nữa.
 - Warning `mel_filterbank` khi khởi động là bình thường — liên quan đến audio encoder của Gemma 4 multimodal, không ảnh hưởng đến text generation.
 - Tốc độ generate phụ thuộc vào phần cứng. Trên Orange Pi 5 với CPU, khoảng 5–15 token/giây.
-- Token/s được đo bằng `engine.tokenize()` trên output thực tế — phản ánh đúng throughput của model, không bao gồm thời gian mạng.
+- Token/s dùng `engine.tokenize()` nếu có, fallback về ước tính `len(text) // 4` nếu không khả dụng.
 - Markdown được render bằng [marked.js](https://marked.js.org/) trực tiếp trên trình duyệt, không qua server.
+- Chỉ dùng bản `-it` (instruction-tuned) cho chat — bản base model không phù hợp hội thoại.
 
 ---
 

server.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import uuid
 import time
 from pathlib import Path
@@ -15,9 +16,71 @@ from contextlib import asynccontextmanager
 
 # ── Config ────────────────────────────────────────────────────────────────────
 
-MODEL_PATH   = "gemma-4-E4B-it.litertlm"
+MODELS_DIR   = Path(__file__).parent / "models"
 TEMPLATE_DIR = Path(__file__).parent / "templates"
 
+AVAILABLE_MODELS = {
+    "gemma-4-E2B-it": {
+        "file": "gemma-4-E2B-it.litertlm",
+        "repo": "google/gemma-4-E2B-it",
+        "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",
+        "desc": "Gemma 4 Edge 4B — thông minh hơn, chậm hơn",
+    },
+}
+
+# ── CLI: chọn model khi khởi động ────────────────────────────────────────────
+
+def select_model() -> Path:
+    print("\n" + "="*52)
+    print("  LiteRT-LM Server — Chọn model")
+    print("="*52)
+
+    for i, (key, info) in enumerate(AVAILABLE_MODELS.items(), 1):
+        model_path = MODELS_DIR / info["file"]
+        status = "✓ có sẵn" if model_path.exists() else "✗ chưa tải"
+        print(f"  [{i}] {key}")
+        print(f"      {info['desc']}")
+        print(f"      {status}")
+        print()
+
+    while True:
+        try:
+            choice = input("Chọn model (1/2): ").strip()
+            idx = int(choice) - 1
+            if 0 <= idx < len(AVAILABLE_MODELS):
+                key = list(AVAILABLE_MODELS.keys())[idx]
+                info = AVAILABLE_MODELS[key]
+                model_path = MODELS_DIR / info["file"]
+
+                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"    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
+                    else:
+                        sys.exit(0)
+
+                print(f"\n  Đã chọn: {key}")
+                print(f"  Path: {model_path}\n")
+                return model_path
+            else:
+                print("  Vui lòng nhập 1 hoặc 2.")
+        except (ValueError, KeyboardInterrupt):
+            print("\n  Thoát.")
+            sys.exit(0)
+
+# Chọn model trước khi FastAPI khởi động
+MODELS_DIR.mkdir(exist_ok=True)
+MODEL_PATH = select_model()
+
 # ── Models ───────────────────────────────────────────────────────────────────
 
 class PromptRequest(BaseModel):
@@ -28,12 +91,23 @@ class PromptRequest(BaseModel):
 ml_models = {}
 sessions: dict = {}  # session_id -> conversation object
 
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+def count_tokens(engine, text: str) -> int:
+    try:
+        return len(engine.tokenize(text))
+    except Exception:
+        return max(1, len(text) // 4)
+
 # ── Lifespan ─────────────────────────────────────────────────────────────────
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.CPU)
+    print(f"  Loading model: {MODEL_PATH.name} ...")
+    engine = litert_lm.Engine(str(MODEL_PATH), backend=litert_lm.Backend.CPU)
     ml_models["engine"] = engine
+    ml_models["model_name"] = MODEL_PATH.stem
+    print(f"  Model ready: {MODEL_PATH.name}\n")
     yield
     sessions.clear()
     del ml_models["engine"]
@@ -42,6 +116,16 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(title="LiteRT-LM API", lifespan=lifespan)
 
+# ── REST: info ────────────────────────────────────────────────────────────────
+
+@app.get("/info")
+async def info():
+    """Return current loaded model info."""
+    return {
+        "model": ml_models.get("model_name", "unknown"),
+        "sessions": len(sessions),
+    }
+
 # ── REST: stateless single-turn ───────────────────────────────────────────────
 
 @app.post("/generate")
@@ -56,7 +140,7 @@ async def generate_text(request: PromptRequest):
         result = conversation.send_message(request.prompt)
         elapsed = time.perf_counter() - t0
         text = result["content"][0]["text"]
-        num_tokens = len(engine.tokenize(text))
+        num_tokens = count_tokens(engine, text)
         tps = round(num_tokens / elapsed, 2) if elapsed > 0 else 0
         return {
             "response": text,
@@ -93,7 +177,7 @@ async def chat(session_id: str, request: PromptRequest):
         result = sessions[session_id].send_message(request.prompt)
         elapsed = time.perf_counter() - t0
         text = result["content"][0]["text"]
-        num_tokens = len(engine.tokenize(text)) if engine else 0
+        num_tokens = count_tokens(engine, text)
         tps = round(num_tokens / elapsed, 2) if elapsed > 0 else 0
         return {
             "session_id": session_id,