Function process_files

AskAgent/python/misc/embedding.py:42–107

(input_file, size="small", model="text-embedding-3-small", num_to_process=10000000)

    return text.encode('utf-8', errors='ignore').decode('utf-8')

def process_files(input_file, size="small", model="text-embedding-3-small", num_to_process=10000000):
    num_done = 0
    if (size == "small"):
        output_path = EMBEDDINGS_PATH_SMALL + "/" + input_file + ".txt"
        model = EMBEDDING_MODEL_SMALL
    elif (size == "compact"):
        output_path = EMBEDDINGS_PATH_COMPACT + "/" + input_file + ".txt"
        model = EMBEDDING_MODEL_SMALL
    else:
        output_path = EMBEDDINGS_PATH_LARGE + "/" + input_file + ".txt"
        model = EMBEDDING_MODEL_LARGE

    input_path = JSONL_PATH + input_file + "_schemas.txt"

    if (size == "compact"):
        input_path = JSONL_PATH_COMPACT + input_file + "_schemas.txt"

    try:
        with open(input_path) as input_file, \
             open(output_path, 'w', encoding='utf-8') as output_file:

            batch = []
            batch_urls = []
            batch_jsons = []

            for line in input_file:
                # Skip empty lines
                if not line.strip():
                    continue

                line = clean_utf8(line)
                try:
                    # Split line by tab
                    url, json_str = line.strip().split('\t')

                    batch_urls.append(url)
                    batch_jsons.append(json_str)
                    batch.append(json_str[0:6000])
                    num_done += 1
                    # Process batch when it reaches size 100
                    if len(batch) == 100 or (num_done > num_to_process):
                        # Get embeddings for the batch
                        embeddings = client.embeddings.create(input=batch, model=model).data

                        # Write results for the batch
                        for i in range(len(batch)):
                            output_file.write(f"{batch_urls[i]}\t{batch_jsons[i]}\t{embeddings[i].embedding}\n")
                        print(f"Processed {num_done} lines")
                        # Clear the batches
                        batch = []
                        batch_urls = []
                        batch_jsons = []
                        time.sleep(5)
                except Exception as e:
                    print(f"Error processing line: {e!s}")
                    continue
                if num_done > num_to_process:
                    break