Spaces:
Runtime error
Runtime error
File size: 1,490 Bytes
1ae7d73 |
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 58 59 60 61 62 63 64 65 66 67 68 |
# %%
import shutil
import os
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
# %%
# loading the riddle data into docs
data_file = "data/riddles_data"
loader = TextLoader(data_file)
docs = loader.load()
# create the text splitter, splitted exactly line-by-line
text_splitter = CharacterTextSplitter(
separator = "\n",
chunk_size = 0,
chunk_overlap = 0,
length_function = len,
is_separator_regex = False,
)
# get the splits
splits = text_splitter.split_documents(docs)
# %%
# loading the vector encoder
model_name = "shibing624/text2vec-base-chinese"
encode_kwargs = {'normalize_embeddings': False}
model_kwargs = {'device': 'cpu'}
huggingface_embeddings= HuggingFaceEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs = encode_kwargs
)
# %%
# vectordb with Chroma
persist_directory = 'chroma/'
# %%
# remove the old one when rebuilt the database
if os.path.exists(persist_directory):
shutil.rmtree(persist_directory) # remove old database files if any
# %%
# load the riddles documents to vectordb
vectordb = Chroma.from_documents(
documents=splits,
embedding=huggingface_embeddings,
persist_directory=persist_directory,
collection_metadata={"hnsw:space": "cosine"}
)
# %%
vectordb.persist()
print(vectordb._collection.count())
# %%
|