import torch, librosa, laion_clap, functools, warnings
warnings.filterwarnings('ignore')

original_load = torch.load
torch.load = functools.partial(original_load, weights_only=False)
original_load_state_dict = torch.nn.Module.load_state_dict

def tolerant_load_state_dict(self, state_dict, strict=True, assign=False):
    return original_load_state_dict(self, state_dict, strict=False, assign=assign)
torch.nn.Module.load_state_dict = tolerant_load_state_dict

model = laion_clap.CLAP_Module(enable_fusion=False)
model.load_ckpt() 

torch.load = original_load
torch.nn.Module.load_state_dict = original_load_state_dict

def get_clap_embedding(path):
    audio_data, _ = librosa.load(path, sr=48000)
    audio_data = audio_data.reshape(1, -1)
    with torch.no_grad():
        audio_embed = model.get_audio_embedding_from_data(x=audio_data)
    return audio_embed.flatten()