Update README.md

README.md

@@ -1,3 +1,60 @@
 Very well working 2 layer NN for projection. 
 Trained with custom N Pairs loss + hard negative mining.
-  margin = 1.0
+  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 2 hidden layers
+        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),  # Second 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!")