code-llama-7b-projection-largev2.11

Very well working 2 layer NN for projection. Trained with custom N Pairs loss + hard negative mining. margin = 1.0

Load in using following code:

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoTokenizer, AutoModel
from huggingface_hub import hf_hub_download
import json
from types import SimpleNamespace

#model architecture - needed since this is a custom model
class ProjectionModel(nn.Module):
    def __init__(self, config):
        super(ProjectionModel, self).__init__()
        self.config = config
        self.c_code_encoder = AutoModel.from_pretrained("microsoft/codebert-base")
        self.pseudocode_encoder = AutoModel.from_pretrained("microsoft/codebert-base")

        #Projection network with 1 hidden layer
        self.projection = nn.Sequential(
            nn.Linear(config.embedding_dim, config.hidden_dim),
            nn.ReLU(),  # First activation function
            nn.Linear(config.hidden_dim, config.hidden_dim),  #Hidden layer
            nn.ReLU(),  # Second activation function
            nn.Linear(config.hidden_dim, config.embedding_dim)  #Output layer projecting back to the original embedding space
        )

    def forward(self, c_code_inputs, pseudocode_inputs):
        #Encode C code and pseudocode
        c_code_embedding = self.c_code_encoder(**c_code_inputs).last_hidden_state.mean(dim=1)
        pseudocode_embedding = self.pseudocode_encoder(**pseudocode_inputs).last_hidden_state.mean(dim=1)

        #Apply the projection network to the pseudocode embeddings
        projected_pseudocode_embedding = self.projection(pseudocode_embedding)

        return c_code_embedding, projected_pseudocode_embedding

model_name = "aircrypto/code-llama-7b-projection-largev2.11"
config_file = hf_hub_download(repo_id=model_name, filename="config.json")

with open(config_file, 'r') as f:
    config_dict = json.load(f)
config = SimpleNamespace(**config_dict)

model = ProjectionModel(config)
model_path = hf_hub_download(repo_id=model_name, filename="pytorch_model.bin")
state_dict = torch.load(model_path, map_location="cpu")
model.load_state_dict(state_dict)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

print("Model loaded successfully!")

tokenizer = AutoTokenizer.from_pretrained("aircrypto/code-llama-7b-projection-largev2.11")

print("Tokenizer loaded successfully!")