diff --git a/.gitignore b/.gitignore index 12b823e..d4b2dfa 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ gemma-4-E2B-it.* -gemma-4-E4B-it.* \ No newline at end of file +gemma-4-E4B-it.* +models/ \ No newline at end of file diff --git a/README.md b/README.md index c5e1bd9..a882658 100644 --- a/README.md +++ b/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 ``` --- @@ -21,30 +28,88 @@ pip install fastapi uvicorn pydantic ``` . -├── app.py # REST API đơn giản, single-turn -├── server.py # REST API đầy đủ + Web UI, multi-turn sessions +├── 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/` -``` -gemma-4-E2B-it.litertlm ← model file -app.py -server.py +Xem lệnh tải ở trên. + +### Bước 2 — Chạy server + +```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. - -### Chạy - -```bash -python server.py -``` - -Server khởi động tại `http://0.0.0.0:8000` +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. --- ### 🌐 Web UI -Mở trình duyệt và truy cập: - -``` -http://<địa-chỉ-ip>:8000 -``` +Mở trình duyệt và truy cập `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.) -- **Đ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ý +- **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` 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) -``` - ---- - -## 🧪 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!"}' +engine = litert_lm.Engine(str(MODEL_PATH), backend=litert_lm.Backend.GPU) ``` --- @@ -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. --- diff --git a/server.py b/server.py index 179fb0d..3a5025b 100644 --- a/server.py +++ b/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,