sem_alloc.py

import os
import backend.processer as processer
import json
import numpy as np
import pickle

MIME_TYPES = {'.jpg': 'jpeg', '.jpeg': 'jpeg', '.png':'png', '.webp':'webp'}
def semantic_allocation(directory_path):
    with open("cold_storage.json", "w") as f:
        f.write("{\n")
        for root, _, files in os.walk(directory_path):
            for file_ in files:
                full_path = os.path.join(root, file_)
                try:
                    _, ext = os.path.splitext(file_)
                    ext = ext.lower()

                    if ext in ['.mp3', '.mp4', '.m4a']:
                        output = processer.process_audio(full_path).replace('\n', ' ')
                        s = fr"{output}"
                        f.write(f"\"{full_path}\":\"{s}\",\n")
                    elif ext in ['.jpg', '.jpeg', '.png', '.webp']:
                        output = processer.process_image(full_path, MIME_TYPES[ext]).replace('\n', ' ')
                        s = fr"{output}"
                        f.write(f"\"{full_path}\":\"{s}\",\n")
                    elif ext in ['.txt']:
                        with open(full_path, 'r', encoding='utf-8') as txt_file:
                            output = txt_file.read().replace('\n', ' ')
                            s = fr"{output}"
                            f.write(f"\"{full_path}\":\"{s}\",\n")
                    else:
                        print(f"Unsupported file type for {full_path}")
                    print(f"Embedding file: {full_path}")

                except Exception as e:
                    print(f"Error embedding {full_path}: {e}")
        f.write("}")


def embed_everything(cold_storage_file):
    with open(cold_storage_file, 'r') as f:
        summaries = json.load(f)
        num_summaries = len(summaries)
        batch_size = 15
        summaries_items = list(summaries.items())
        embeddings = {}
        embedding_to_file = {}
        embeddings_ = []

        for i in range(0, num_summaries, batch_size):
            batch = summaries_items[i:i + batch_size]
            batch_texts = [text for _, text in batch]

            batch_embeddings_response = processer.gemini_embed(strings=batch_texts)
            batch_embeddings = batch_embeddings_response.embeddings
            
            count = 0
            for (file_path, _), embedding_data in zip(batch, batch_embeddings):
                embedding_array = np.array(embedding_data.values)
                embeddings[file_path] = {"array": embedding_array.tolist(), "id": i + count}
                embeddings_.append(embedding_array)
                count += 1


            print(f"Processed batch {i // batch_size + 1}/{(num_summaries + batch_size - 1) // batch_size}")

    with open("embeddings.json", "w") as f2:
        json.dump(embeddings, f2)
    with open("embeddings.pkl", "wb") as pkl_file:
        embeddings_ = np.asarray(embeddings_)
        # np.savez_compressed('embeddings.npz', embeddings=embeddings_ )
        np.save('embeddings.npy', embeddings_)
        pickle.dump(embeddings_, pkl_file)

    return embeddings, embeddings_


# directory_path = "/Users/derekzhu/Code/EmbedAnything/EmbedAnything/test"
# semantic_allocation(directory_path)
embed_everything("/Users/derekzhu/Code/EmbedAnything/EmbedAnything/cold_storage.json")