mirror of
https://github.com/HKUDS/LightRAG.git
synced 2025-06-26 22:00:19 +00:00
231 lines
7.5 KiB
Python
231 lines
7.5 KiB
Python
# pip install -q -U google-genai to use gemini as a client
|
|
|
|
import os
|
|
from typing import Optional
|
|
import dataclasses
|
|
from pathlib import Path
|
|
import hashlib
|
|
import numpy as np
|
|
from google import genai
|
|
from google.genai import types
|
|
from dotenv import load_dotenv
|
|
from lightrag.utils import EmbeddingFunc, Tokenizer
|
|
from lightrag import LightRAG, QueryParam
|
|
from sentence_transformers import SentenceTransformer
|
|
from lightrag.kg.shared_storage import initialize_pipeline_status
|
|
import sentencepiece as spm
|
|
import requests
|
|
|
|
import asyncio
|
|
import nest_asyncio
|
|
|
|
# Apply nest_asyncio to solve event loop issues
|
|
nest_asyncio.apply()
|
|
|
|
load_dotenv()
|
|
gemini_api_key = os.getenv("GEMINI_API_KEY")
|
|
|
|
WORKING_DIR = "./dickens"
|
|
|
|
if os.path.exists(WORKING_DIR):
|
|
import shutil
|
|
|
|
shutil.rmtree(WORKING_DIR)
|
|
|
|
os.mkdir(WORKING_DIR)
|
|
|
|
|
|
class GemmaTokenizer(Tokenizer):
|
|
# adapted from google-cloud-aiplatform[tokenization]
|
|
|
|
@dataclasses.dataclass(frozen=True)
|
|
class _TokenizerConfig:
|
|
tokenizer_model_url: str
|
|
tokenizer_model_hash: str
|
|
|
|
_TOKENIZERS = {
|
|
"google/gemma2": _TokenizerConfig(
|
|
tokenizer_model_url="https://raw.githubusercontent.com/google/gemma_pytorch/33b652c465537c6158f9a472ea5700e5e770ad3f/tokenizer/tokenizer.model",
|
|
tokenizer_model_hash="61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2",
|
|
),
|
|
"google/gemma3": _TokenizerConfig(
|
|
tokenizer_model_url="https://raw.githubusercontent.com/google/gemma_pytorch/cb7c0152a369e43908e769eb09e1ce6043afe084/tokenizer/gemma3_cleaned_262144_v2.spiece.model",
|
|
tokenizer_model_hash="1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c",
|
|
),
|
|
}
|
|
|
|
def __init__(
|
|
self, model_name: str = "gemini-2.0-flash", tokenizer_dir: Optional[str] = None
|
|
):
|
|
# https://github.com/google/gemma_pytorch/tree/main/tokenizer
|
|
if "1.5" in model_name or "1.0" in model_name:
|
|
# up to gemini 1.5 gemma2 is a comparable local tokenizer
|
|
# https://github.com/googleapis/python-aiplatform/blob/main/vertexai/tokenization/_tokenizer_loading.py
|
|
tokenizer_name = "google/gemma2"
|
|
else:
|
|
# for gemini > 2.0 gemma3 was used
|
|
tokenizer_name = "google/gemma3"
|
|
|
|
file_url = self._TOKENIZERS[tokenizer_name].tokenizer_model_url
|
|
tokenizer_model_name = file_url.rsplit("/", 1)[1]
|
|
expected_hash = self._TOKENIZERS[tokenizer_name].tokenizer_model_hash
|
|
|
|
tokenizer_dir = Path(tokenizer_dir)
|
|
if tokenizer_dir.is_dir():
|
|
file_path = tokenizer_dir / tokenizer_model_name
|
|
model_data = self._maybe_load_from_cache(
|
|
file_path=file_path, expected_hash=expected_hash
|
|
)
|
|
else:
|
|
model_data = None
|
|
if not model_data:
|
|
model_data = self._load_from_url(
|
|
file_url=file_url, expected_hash=expected_hash
|
|
)
|
|
self.save_tokenizer_to_cache(cache_path=file_path, model_data=model_data)
|
|
|
|
tokenizer = spm.SentencePieceProcessor()
|
|
tokenizer.LoadFromSerializedProto(model_data)
|
|
super().__init__(model_name=model_name, tokenizer=tokenizer)
|
|
|
|
def _is_valid_model(self, model_data: bytes, expected_hash: str) -> bool:
|
|
"""Returns true if the content is valid by checking the hash."""
|
|
return hashlib.sha256(model_data).hexdigest() == expected_hash
|
|
|
|
def _maybe_load_from_cache(self, file_path: Path, expected_hash: str) -> bytes:
|
|
"""Loads the model data from the cache path."""
|
|
if not file_path.is_file():
|
|
return
|
|
with open(file_path, "rb") as f:
|
|
content = f.read()
|
|
if self._is_valid_model(model_data=content, expected_hash=expected_hash):
|
|
return content
|
|
|
|
# Cached file corrupted.
|
|
self._maybe_remove_file(file_path)
|
|
|
|
def _load_from_url(self, file_url: str, expected_hash: str) -> bytes:
|
|
"""Loads model bytes from the given file url."""
|
|
resp = requests.get(file_url)
|
|
resp.raise_for_status()
|
|
content = resp.content
|
|
|
|
if not self._is_valid_model(model_data=content, expected_hash=expected_hash):
|
|
actual_hash = hashlib.sha256(content).hexdigest()
|
|
raise ValueError(
|
|
f"Downloaded model file is corrupted."
|
|
f" Expected hash {expected_hash}. Got file hash {actual_hash}."
|
|
)
|
|
return content
|
|
|
|
@staticmethod
|
|
def save_tokenizer_to_cache(cache_path: Path, model_data: bytes) -> None:
|
|
"""Saves the model data to the cache path."""
|
|
try:
|
|
if not cache_path.is_file():
|
|
cache_dir = cache_path.parent
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
with open(cache_path, "wb") as f:
|
|
f.write(model_data)
|
|
except OSError:
|
|
# Don't raise if we cannot write file.
|
|
pass
|
|
|
|
@staticmethod
|
|
def _maybe_remove_file(file_path: Path) -> None:
|
|
"""Removes the file if exists."""
|
|
if not file_path.is_file():
|
|
return
|
|
try:
|
|
file_path.unlink()
|
|
except OSError:
|
|
# Don't raise if we cannot remove file.
|
|
pass
|
|
|
|
# def encode(self, content: str) -> list[int]:
|
|
# return self.tokenizer.encode(content)
|
|
|
|
# def decode(self, tokens: list[int]) -> str:
|
|
# return self.tokenizer.decode(tokens)
|
|
|
|
|
|
async def llm_model_func(
|
|
prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
|
|
) -> str:
|
|
# 1. Initialize the GenAI Client with your Gemini API Key
|
|
client = genai.Client(api_key=gemini_api_key)
|
|
|
|
# 2. Combine prompts: system prompt, history, and user prompt
|
|
if history_messages is None:
|
|
history_messages = []
|
|
|
|
combined_prompt = ""
|
|
if system_prompt:
|
|
combined_prompt += f"{system_prompt}\n"
|
|
|
|
for msg in history_messages:
|
|
# Each msg is expected to be a dict: {"role": "...", "content": "..."}
|
|
combined_prompt += f"{msg['role']}: {msg['content']}\n"
|
|
|
|
# Finally, add the new user prompt
|
|
combined_prompt += f"user: {prompt}"
|
|
|
|
# 3. Call the Gemini model
|
|
response = client.models.generate_content(
|
|
model="gemini-1.5-flash",
|
|
contents=[combined_prompt],
|
|
config=types.GenerateContentConfig(max_output_tokens=500, temperature=0.1),
|
|
)
|
|
|
|
# 4. Return the response text
|
|
return response.text
|
|
|
|
|
|
async def embedding_func(texts: list[str]) -> np.ndarray:
|
|
model = SentenceTransformer("all-MiniLM-L6-v2")
|
|
embeddings = model.encode(texts, convert_to_numpy=True)
|
|
return embeddings
|
|
|
|
|
|
async def initialize_rag():
|
|
rag = LightRAG(
|
|
working_dir=WORKING_DIR,
|
|
# tiktoken_model_name="gpt-4o-mini",
|
|
tokenizer=GemmaTokenizer(
|
|
tokenizer_dir=(Path(WORKING_DIR) / "vertexai_tokenizer_model"),
|
|
model_name="gemini-2.0-flash",
|
|
),
|
|
llm_model_func=llm_model_func,
|
|
embedding_func=EmbeddingFunc(
|
|
embedding_dim=384,
|
|
max_token_size=8192,
|
|
func=embedding_func,
|
|
),
|
|
)
|
|
|
|
await rag.initialize_storages()
|
|
await initialize_pipeline_status()
|
|
|
|
return rag
|
|
|
|
|
|
def main():
|
|
# Initialize RAG instance
|
|
rag = asyncio.run(initialize_rag())
|
|
file_path = "story.txt"
|
|
with open(file_path, "r") as file:
|
|
text = file.read()
|
|
|
|
rag.insert(text)
|
|
|
|
response = rag.query(
|
|
query="What is the main theme of the story?",
|
|
param=QueryParam(mode="hybrid", top_k=5, response_type="single line"),
|
|
)
|
|
|
|
print(response)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|