Публикация решения для taskId 6a1d75c5fd30e81cf3126ae7: update src/utils.py
This commit is contained in:
+8
-3
@@ -9,6 +9,9 @@ from langchain_ollama import OllamaEmbeddings
|
|||||||
# Chroma collection name
|
# Chroma collection name
|
||||||
COLLECTION_NAME = "faq"
|
COLLECTION_NAME = "faq"
|
||||||
|
|
||||||
|
# Persistence directory from config
|
||||||
|
from src.config import CHROMA_PERSIST_DIR
|
||||||
|
|
||||||
|
|
||||||
def load_faq_to_chroma(md_dir: str = "data") -> None:
|
def load_faq_to_chroma(md_dir: str = "data") -> None:
|
||||||
"""Load all .md files from md_dir into a Chroma vector store.
|
"""Load all .md files from md_dir into a Chroma vector store.
|
||||||
@@ -22,21 +25,23 @@ def load_faq_to_chroma(md_dir: str = "data") -> None:
|
|||||||
all_docs.extend(docs)
|
all_docs.extend(docs)
|
||||||
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||||
docs = splitter.split_documents(all_docs)
|
docs = splitter.split_documents(all_docs)
|
||||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
embeddings = OllamaEmbeddings("nomic-embed-text")
|
||||||
# Create or update Chroma collection
|
# Create or update Chroma collection with persistence
|
||||||
Chroma.from_documents(
|
Chroma.from_documents(
|
||||||
docs,
|
docs,
|
||||||
embeddings,
|
embeddings,
|
||||||
collection_name=COLLECTION_NAME,
|
collection_name=COLLECTION_NAME,
|
||||||
|
persist_directory=CHROMA_PERSIST_DIR,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def search_course_docs(query: str, k: int = 3) -> List[str]:
|
def search_course_docs(query: str, k: int = 3) -> List[str]:
|
||||||
"""Return top k document snippets from the persisted Chroma store."""
|
"""Return top k document snippets from the persisted Chroma store."""
|
||||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
embeddings = OllamaEmbeddings("nomic-embed-text")
|
||||||
store = Chroma(
|
store = Chroma(
|
||||||
collection_name=COLLECTION_NAME,
|
collection_name=COLLECTION_NAME,
|
||||||
embedding_function=embeddings,
|
embedding_function=embeddings,
|
||||||
|
persist_directory=CHROMA_PERSIST_DIR,
|
||||||
)
|
)
|
||||||
results = store.similarity_search(query, k=k)
|
results = store.similarity_search(query, k=k)
|
||||||
return [doc.page_content for doc in results]
|
return [doc.page_content for doc in results]
|
||||||
|
|||||||
Reference in New Issue
Block a user