Spaces:
Running
Running
File size: 2,009 Bytes
f126864 797c083 f126864 |
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 |
import os
import argparse
from typing import Dict
import dotenv
from pathlib import Path
from tqdm import tqdm
from pymilvus import MilvusClient, model
_ = dotenv.load_dotenv()
def create_collection(client: MilvusClient, collection_name: str, dimension: int):
if client.has_collection(collection_name=collection_name):
client.drop_collection(collection_name=collection_name)
client.create_collection(
collection_name=collection_name,
dimension=dimension,
)
def main(args: Dict):
client = MilvusClient("milvus.db")
embedding_fn = model.dense.OpenAIEmbeddingFunction(
model_name=args.model_name,
api_key=os.environ.get('OPENAI_API_KEY'),
dimensions=args.dimension
)
create_collection(client, args.collection_name, args.dimension)
docs = Path(args.docs_dir)
md_file_paths = list(docs.rglob('*.md'))
mdx_file_paths = list(docs.rglob('*.mdx'))
all_file_paths = md_file_paths + mdx_file_paths
docs, payloads = [], []
for file in tqdm(all_file_paths):
embed_string = str(file).replace('docs/', '').replace('.mdx', '').replace('.md', '').replace('/', ' ')
docs.append(embed_string)
payloads.append({'file_path': str(file)})
vectors = embedding_fn.encode_documents(docs)
data = [
{"id": i, "vector": vectors[i], "text": docs[i], **payloads[i]}
for i in range(len(vectors))
]
response = client.insert(collection_name=args.collection_name, data=data)
print(f"Inserted {response['insert_count']} vectors into collection {args.collection_name}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--collection_name", type=str, default="hf_docs")
parser.add_argument("--model_name", type=str, default="text-embedding-3-small")
parser.add_argument("--dimension", type=int, default=1536)
parser.add_argument("--docs_dir", type=str, default="docs")
args = parser.parse_args()
main(args) |