data_preparation.py

import os
import argparse
from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
import chromadb
from langchain.vectorstores import Chroma


def parse_args(args=None):
    parser = argparse.ArgumentParser()

    parser.add_argument('--vectordb_name', type=str, default=None, help="the name of vectordb")

    return parser.parse_args(args)


if __name__ == "__main__":
    args = parse_args()
    embeddings = SentenceTransformerEmbeddings(
        model_name=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'inference/all-MiniLM-L6-v2'),
        model_kwargs={"device": "cuda"})
    chroma = chromadb.PersistentClient(
        path=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'database/', args.vectordb_name))

    print("-----------MQuAKE Preparation----------")
    collection = chroma.create_collection("wiki_mquake_2104")

    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data/mquake_2104/')
    files = os.listdir(path)
    cnt = 1
    for file in files:
        with open(path + file, 'r') as f:
            raw_text = f.read()
            collection.add(
                documents=[raw_text],
                metadatas=[{"title": file, "source": "wiki"}],
                ids=["id_" + file])
            print(cnt, file, len(raw_text.split()))
        cnt += 1

    print("-----------FreshQA Preparation----------")
    collection = chroma.create_collection("wiki_freshqa_2104")

    path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data/freshqa_2104/')
    files = os.listdir(path)
    cnt = 1
    for file in files:
        with open(path + file, 'r') as f:
            raw_text = f.read()
            collection.add(
                documents=[raw_text],
                metadatas=[{"title": file, "source": "wiki"}],
                ids=["id_" + file])
            print(cnt, file, len(raw_text.split()))
        cnt += 1