[apps] Embedding Gemma2 RAG on edge

apps/RAG-on-Edge/README.md

@@ -0,0 +1,53 @@
+# RAG on Edge
+
+End-to-end Retrieval-Augmented Generation running fully on-device using
+[Qdrant Edge](https://qdrant.tech/) for vector search and
+[LiteRT-LM](https://github.com/google-ai-edge/litert-lm) for embeddings and
+generation — no network calls at inference time.
+
+The pipeline indexes the **Gemma 4 technical report** PDF and answers questions
+about it using EmbeddingGemma2 (retrieval) and Gemma 4 E2B (generation).
+
+## Components
+
+| File | Purpose |
+|------|---------|
+| `download_models.py` | Downloads the two `.litertlm` model files from Hugging Face. |
+| `index.py` | Parses the PDF, chunks it, embeds the chunks, and stores them in a local Qdrant Edge shard (`db_gemma`). |
+| `app.py` | Loads the shard + models and answers a set of queries via retrieval + Gemma 4 generation. |
+| `quiet_native.py` | Silences noisy native (C++) stderr output from the runtime. |
+
+## Models
+
+- **Embedding:** `litert-community/embeddinggemma-2-740m-litert-lm` → `models/embeddinggemma2/embeddinggemma-2-740m.litertlm` (768-dim)
+- **Generation:** `litert-community/gemma-4-E2B-it-litert-lm` → `models/gemma4/gemma-4-E2B-it.litertlm`
+
+Only the single required `.litertlm` file is pulled per model.
+
+## Setup
+
+```bash
+# install dependencies
+pip install -e .
+
+# download the embedding + generation models
+python download_models.py          # or: runedge-download-model
+```
+
+## Usage
+
+```bash
+# 1. build the vector index from the PDF (../sample-data/gemma4_technical_report.pdf)
+python index.py                    # or: runedge-index
+
+# 2. run retrieval + generation over the sample queries
+python app.py
+```
+
+`index.py` creates/loads the shard at `db_gemma/`; if the shard already exists it
+is reused rather than rebuilt. Both scripts cache runtime artifacts in `.cache/`.
+
+## Requirements
+
+- Python >= 3.11
+- A GPU-capable backend (both the embedding and generation engines use `Backend.GPU()`)

apps/RAG-on-Edge/app.py

@@ -0,0 +1,97 @@
+import os
+import litert_lm
+from litert_lm import Backend, EmbeddingEngine
+from qdrant_edge import EdgeShard, Query, QueryRequest
+from quiet_native import silence_stderr
+os.environ["HF_HUB_OFFLINE"] = "1"
+silence_stderr()
+
+EMBED_MODEL_NAME = "models/embeddinggemma2/embeddinggemma-2-740m.litertlm"
+SHARD_PATH = "db_gemma"
+GEMMA4_LITELLM = "models/gemma4/gemma-4-E2B-it.litertlm"
+CACHE_DIR = ".cache"
+
+class GemmaInference:
+    def __init__(
+        self,
+        embed_model_name: str = EMBED_MODEL_NAME,
+        shard_path: str = SHARD_PATH,
+        gemma_model_path: str = GEMMA4_LITELLM,
+    ) -> None:
+
+        os.makedirs(CACHE_DIR, exist_ok=True)
+        self.model = EmbeddingEngine(
+            embed_model_name, backend=Backend.GPU(), cache_dir=CACHE_DIR
+        )
+
+        self.shard = EdgeShard.load(shard_path)
+
+        self.gemma_model_path = gemma_model_path
+        self.engine = litert_lm.Engine(
+            self.gemma_model_path,
+            backend=litert_lm.Backend.GPU(),
+        )
+
+    def close(self) -> None:
+        self.engine.close()
+        self.model.close()
+
+    def retrieve_context(self, user_query: str, limit: int = 2) -> str:
+        query_vector = self.model.compute_embedding(user_query).embedding
+        results = self.shard.query(
+            QueryRequest(
+                query=Query.Nearest(query_vector),
+                limit=limit,
+                with_payload=True,
+            )
+        )
+        return "".join(point.payload["text"] for point in results)
+
+    def _system_messages(self) -> list[litert_lm.Message]:
+        return [
+            litert_lm.Message.system(
+                "You are a Gemma 4 technical report assistant. "
+                "Answer the user question using only the provided context. "
+                "If the answer is not in the context, say: I don't know"
+                "Do not guess or add extra explanation."
+                "Answer in 1-2 short sentences. Do not elaborate unless told by user"
+            )
+        ]
+
+    def answer_stream(self, user_query: str, limit: int = 3, context: str | None = None):
+        if context is None:
+            context = self.retrieve_context(user_query, limit=limit)
+        with self.engine.create_conversation(
+            messages=self._system_messages(),
+            sampler_config=litert_lm.SamplerConfig(top_k=1),
+        ) as conversation:
+            for chunk in conversation.send_message_async(
+                f"Context:\n{context}\n\nQuestion:\n{user_query}"
+            ):
+                for item in chunk.get("content", []):
+                    if item.get("type") == "text" and item.get("text"):
+                        yield item["text"]
+
+    def answer(self, user_query: str, limit: int = 3) -> str:
+        return "".join(self.answer_stream(user_query, limit=limit))
+
+QUERIES = [
+    "What is the parameter range of models in the Gemma 4 suite?",
+    "What architecture is Gemma 4 based on?",
+    "What context length does Gemma 4 support?",
+    "How does Gemma 4 perform on benchmarks compared to prior Gemma models?"
+]
+
+def main() -> None:
+    inference = GemmaInference()
+
+    for i, user_query in enumerate(QUERIES, start=1):
+        print(f"\n[{i}] {user_query}")
+        context = inference.retrieve_context(user_query)
+        answer = "".join(inference.answer_stream(user_query, context=context))
+        print(answer)
+
+    inference.close()
+
+if __name__ == "__main__":
+    main()

apps/RAG-on-Edge/download_models.py

@@ -0,0 +1,35 @@
+from huggingface_hub import hf_hub_download
+
+def download_embedding_model() -> None:
+    print("Downloading EmbeddingGemma2 LiteRT model...")
+    EMBEDDING_REPO_ID = "litert-community/embeddinggemma-2-740m-litert-lm"
+    EMBEDDING_FILENAME = "embeddinggemma-2-740m.litertlm"
+    EMBEDDING_LOCAL_DIR = "models/embeddinggemma2"
+
+    model_path = hf_hub_download(
+        repo_id=EMBEDDING_REPO_ID,
+        filename=EMBEDDING_FILENAME,
+        local_dir=EMBEDDING_LOCAL_DIR,
+    )
+
+    print(f"EmbeddingGemma2 model saved at: {model_path}")
+
+def download_gemma_litert_model() -> None:
+    print("Downloading Gemma LiteRT model...")
+    GEMMA_REPO_ID = "litert-community/gemma-4-E2B-it-litert-lm"
+    GEMMA_FILENAME = "gemma-4-E2B-it.litertlm"
+    GEMMA_LOCAL_DIR = "models/gemma4"
+
+    model_path = hf_hub_download(
+        repo_id=GEMMA_REPO_ID,
+        filename=GEMMA_FILENAME,
+        local_dir=GEMMA_LOCAL_DIR,
+    )
+    print(f"Gemma LiteRT model saved at: {model_path}")
+
+def main() -> None:
+    download_embedding_model()
+    download_gemma_litert_model()
+
+if __name__ == "__main__":
+    main()

apps/RAG-on-Edge/index.py

@@ -0,0 +1,115 @@
+from pathlib import Path
+
+import edgeparse
+from qdrant_edge import (
+    Distance,EdgeConfig,EdgeShard,EdgeVectorParams,
+    Point,Query,QueryRequest,UpdateOperation,
+)
+from litert_lm import Backend, EmbeddingEngine
+
+from quiet_native import silence_stderr
+silence_stderr()
+
+MODEL_DIM = 768  
+EMBED_MODEL_PATH = "models/embeddinggemma2/embeddinggemma-2-740m.litertlm"
+DOCUMENT_PATH = "../sample-data/gemma4_technical_report.pdf"
+SHARD_PATH = "db_gemma"
+CACHE_DIR = ".cache"
+
+class DocumentIngestion:
+    def __init__(
+        self,
+        document_path: str,
+        model_path: str,
+        chunk_size: int = 1024,
+        overlap: int = 0,
+    ) -> None:
+        self.document_path = Path(document_path)
+        self.shard_path = Path(SHARD_PATH)
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+        Path(CACHE_DIR).mkdir(parents=True, exist_ok=True)
+        self.engine = EmbeddingEngine(
+            model_path, backend=Backend.GPU(), cache_dir=CACHE_DIR
+        )
+
+        self.shard = self._load_or_create_shard()
+
+    def close(self) -> None:
+        self.engine.close()
+
+    def _load_or_create_shard(self) -> EdgeShard:
+        config = EdgeConfig(
+            vectors=EdgeVectorParams(size=MODEL_DIM, distance=Distance.Cosine)
+        )
+
+        self.shard_path.mkdir(parents=True, exist_ok=True)
+
+        if any(self.shard_path.iterdir()):
+            print(f"Already exists {self.shard_path}")
+            return EdgeShard.load(str(self.shard_path))
+        return EdgeShard.create(str(self.shard_path), config)
+
+    def chunk_text(self, text: str) -> list[str]:
+        return [
+            chunk
+            for start in range(0, len(text), self.chunk_size - self.overlap)
+            if (chunk := text[start : start + self.chunk_size].strip())
+        ]
+
+    def index(self) -> int:
+        markdown = edgeparse.convert(str(self.document_path), format="markdown") # load the data
+        chunks = self.chunk_text(markdown)
+
+        responses = self.engine.compute_embedding_batch(chunks)
+        embeddings = [response.embedding for response in responses]
+
+        self.shard.update(
+            UpdateOperation.upsert_points(
+                [
+                    Point(
+                        index,
+                        embeddings[index - 1],
+                        {
+                            "text": chunk,
+                            "source": self.document_path.name,
+                            "chunk_index": index - 1,
+                        },
+                    )
+                    for index, chunk in enumerate(chunks, start=1)
+                ]
+            )
+        )
+        self.shard.optimize()
+        return len(chunks)
+
+    def inference(self, user_query: str, limit: int = 2) -> str:
+        query_vector = self.engine.compute_embedding(user_query).embedding
+        results = self.shard.query(
+            QueryRequest(
+                query=Query.Nearest(query_vector),
+                limit=limit, with_payload=True,
+            )
+        )
+        context = ""
+        for point in results:
+            context += point.payload['text']
+
+        return context
+
+def main() -> None:
+    ingestion = DocumentIngestion(
+        document_path=DOCUMENT_PATH, model_path=EMBED_MODEL_PATH
+    )
+    indexed_count = ingestion.index()
+    print(f"Indexed {indexed_count} chunks")
+
+    user_query = "What is the parameter range of models in the Gemma 4 suite?"
+    context = ingestion.inference(user_query)
+    print(context)
+
+    ingestion.close()
+
+if __name__ == "__main__":
+    main()

apps/RAG-on-Edge/pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "runedge"
+version = "0.1.0"
+description = "End to end RAG application on edge using Qdrant Edge and LiteRT"
+requires-python = ">=3.11"
+dependencies = [
+    "edgeparse==0.2.5",
+    "huggingface-hub==1.19.0",
+    "litert-lm-api==0.18.0",
+    "psutil==7.1.3",
+    "qdrant-edge-py==0.7.2",
+    "streamlit==1.58.0",
+]
+
+[project.scripts]
+runedge-download-model = "download_models:main"
+runedge-index = "index:main"
+
+[tool.setuptools]
+py-modules = [
+    "app",
+    "download_models",
+    "index",
+    "quiet_native",
+]

apps/RAG-on-Edge/quiet_native.py

@@ -0,0 +1,19 @@
+import contextlib
+import os
+
+@contextlib.contextmanager
+def suppress_stderr():
+    stderr_fd = 2
+    saved_stderr_fd = os.dup(stderr_fd)
+    try:
+        with open(os.devnull, "w") as devnull:
+            os.dup2(devnull.fileno(), stderr_fd)
+            yield
+    finally:
+        os.dup2(saved_stderr_fd, stderr_fd)
+        os.close(saved_stderr_fd)
+
+def silence_stderr():
+    devnull_fd = os.open(os.devnull, os.O_WRONLY)
+    os.dup2(devnull_fd, 2)
+    os.close(devnull_fd)

apps/README.md

@@ -8,3 +8,4 @@ Here you'll find practical examples and ready-to-run code to help you bring AI-p
 | [Agentic Security (HDP)](Gemma_4_HDP_Agentic_Security/) | Use Human Delegation Provenance (HDP) to cryptographically verify Gemma's native function-calls. |
 | [Gemma on Cloud Run](Gemma-on-Cloudrun/) | Build and deploy a Gemma-powered application on Cloud Run using the Gemini API. |
 | [Antigravity Hybrid Orchestrator](antigravity-hybrid-orchestrator/) | Hybrid cloud/on-device security auditing pairing Gemini 3.8 Flash planning with a local Gemma 4 swarm via the Google Antigravity SDK. |
+| [RAG on Edge](RAG-on-Edge/) | End-to-end Retrieval-Augmented Generation running fully on-device with Qdrant Edge for vector search and LiteRT-LM (EmbeddingGemma2 + Gemma 4) for embeddings and generation — no network calls at inference time. |