mirror of
https://github.com/lllin000/PaperForge.git
synced 2026-07-22 06:50:53 +00:00
37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from paperforge.embedding._chroma import get_collection
|
|
from paperforge.embedding.providers.openai_compatible import OpenAICompatibleProvider
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def embed_paper(vault: Path, zotero_key: str, chunks: list[dict]) -> int:
|
|
"""Embed chunks for one paper using API and insert into ChromaDB. Returns count."""
|
|
collection = get_collection(vault)
|
|
provider = OpenAICompatibleProvider(vault)
|
|
|
|
texts = [c["text"] for c in chunks]
|
|
ids = [f"{zotero_key}_{c['chunk_index']}" for c in chunks]
|
|
metadatas = [
|
|
{
|
|
"paper_id": zotero_key,
|
|
"section": c["section"],
|
|
"page_number": c["page_number"],
|
|
"chunk_index": c["chunk_index"],
|
|
"token_estimate": c["token_estimate"],
|
|
}
|
|
for c in chunks
|
|
]
|
|
|
|
embeddings = provider.encode(texts)
|
|
collection.add(
|
|
ids=ids,
|
|
embeddings=embeddings,
|
|
documents=texts,
|
|
metadatas=metadatas,
|
|
)
|
|
return len(chunks)
|