Spaces:
Sleeping
Sleeping
File size: 2,049 Bytes
86a4642 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from src.config import CHUNK_SIZE, CHUNK_OVERLAP, EMBEDDING_MODEL
from src.persistence import get_vectorstore_path, save_paper_metadata, load_paper_metadata
import uuid
def chunk_text(md_text: str) -> list[str]:
splitter = RecursiveCharacterTextSplitter(
chunk_size=CHUNK_SIZE,
chunk_overlap=CHUNK_OVERLAP,
separators=["\n\n", "\n", " ", ""]
)
return splitter.split_text(md_text)
def get_embeddings():
"""Get embeddings with GPU fallback to CPU."""
try:
return HuggingFaceEmbeddings(
model_name=EMBEDDING_MODEL,
model_kwargs={"device": "cuda"},
)
except Exception:
print("⚠️ GPU not available, falling back to CPU")
return HuggingFaceEmbeddings(
model_name=EMBEDDING_MODEL,
model_kwargs={"device": "cpu"},
)
def build_vectorstore(chunks: list[str], paper_id: str = None) -> tuple[FAISS, str]:
"""Build and save vectorstore. Returns (vectorstore, paper_id)."""
if paper_id is None:
paper_id = str(uuid.uuid4())[:8]
embeddings = get_embeddings()
vectorstore = FAISS.from_texts(chunks, embeddings)
# Save to disk
vectorstore_path = get_vectorstore_path(paper_id)
vectorstore.save_local(str(vectorstore_path))
return vectorstore, paper_id
def load_vectorstore(paper_id: str) -> FAISS:
"""Load a previously saved vectorstore."""
vectorstore_path = get_vectorstore_path(paper_id)
embeddings = get_embeddings()
if vectorstore_path.exists():
vectorstore = FAISS.load_local(
str(vectorstore_path),
embeddings,
allow_dangerous_deserialization=True
)
return vectorstore
raise FileNotFoundError(f"Vectorstore for paper {paper_id} not found") |