rag-vietnamese / src /embedder.py
thaidinhz1's picture
Week 1: minimal RAG pipeline - loader, chunker, embedder, vector store, CLI
02ce710
Raw History Blame Contribute Delete
655 Bytes
import os
from google import genai
from dotenv import load_dotenv
load_dotenv()
client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
EMBED_MODEL = "gemini-embedding-001"
def embed_texts(texts: list[str]) -> list[list[float]]:
"""Embed list văn bản, trả về list vectors."""
result = client.models.embed_content(
model=EMBED_MODEL,
contents=texts,
)
return [e.values for e in result.embeddings]
def embed_query(query: str) -> list[float]:
"""Embed câu hỏi."""
result = client.models.embed_content(
model=EMBED_MODEL,
contents=query,
)
return result.embeddings[0].values