File size: 2,442 Bytes
2822468 92d9c4e 2822468 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
# Buggy-to-Fixed-Code-ViT1D
A 1D Vision Transformer that maps buggy code embeddings to fixed code embeddings.<br>
Datasets that the model can use can be found on: https://huggingface.co/datasets/ASSERT-KTH/RunBugRun-Final<br>
More details about the models training at: https://github.com/ASSERT-KTH/code-embedding-difference
## Usage
```python
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
import pickle
# Define model architecture
class ViT1D(nn.Module):
def __init__(self, input_size=1024, patch_size=16, emb_dim=256, depth=4, heads=8, mlp_ratio=4):
super().__init__()
assert input_size % patch_size == 0
self.num_patches = input_size // patch_size
self.patch_size = patch_size
self.patch_embed = nn.Linear(patch_size, emb_dim)
self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, emb_dim))
encoder_layer = nn.TransformerEncoderLayer(
d_model=emb_dim, nhead=heads,
dim_feedforward=emb_dim * mlp_ratio, batch_first=True
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=depth)
self.output_layer = nn.Linear(emb_dim * self.num_patches, input_size)
def forward(self, x):
bsz = x.size(0)
x = x.view(bsz, self.num_patches, self.patch_size)
x = self.patch_embed(x) + self.pos_embed
x = self.transformer(x)
x = x.flatten(1)
return self.output_layer(x)
# Load model
model = ViT1D(input_size=1024, patch_size=16, emb_dim=256, depth=4, heads=8, mlp_ratio=4)
model_path = hf_hub_download(
repo_id="ASSERT-KTH/Buggy-to-Fixed-Code-ViT1D",
filename="pytorch_model.pth"
)
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()
# Load and predict
file_path = hf_hub_download(
repo_id="ASSERT-KTH/RunBugRun-Final",
filename="Embeddings_RBR/buggy_fixed_embeddings/buggy_fixed_embeddings_chunk_0000.pkl",
repo_type="dataset"
)
with open(file_path, 'rb') as f:
data = pickle.load(f)
buggy_embeddings = data['buggy_embeddings']
with torch.no_grad():
buggy_tensor = torch.tensor(buggy_embeddings[0:1], dtype=torch.float32)
predicted_fixed = model(buggy_tensor).numpy()
print("Predicted fixed embedding:")
print(predicted_fixed[0])
```
## Model Details
- **Architecture:** 1D Vision Transformer
- **Input:** Buggy code embeddings
- **Output:** Fixed code embeddings |