import torch
import torch.nn as nn

vocab_size     = 80
embed_dim      = 128
num_heads      = 4
num_layers     = 3
context_length = 64

class ResumeEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_emb    = nn.Embedding(vocab_size, embed_dim)
        self.position_emb = nn.Embedding(context_length, embed_dim)
        self.blocks = nn.Sequential(*[
            nn.TransformerEncoderLayer(
                d_model=embed_dim,
                nhead=num_heads,
                dim_feedforward=embed_dim * 4,
                dropout=0.1,
                batch_first=True
            ) for _ in range(num_layers)
        ])
        self.output_head = nn.Linear(embed_dim, vocab_size)

    def forward(self, x):
        B, T = x.shape
        tok = self.token_emb(x)
        pos = self.position_emb(torch.arange(T, device=x.device))
        x   = self.blocks(tok + pos)
        return self.output_head(x)

    def encode(self, x):
        # Returns single vector for a sequence — used for similarity search
        tok = self.token_emb(x)
        pos = self.position_emb(torch.arange(x.shape[-1], device=x.device))
        out = self.blocks(tok + pos)
        return out.mean(dim=-2)

if __name__ == "__main__":
    model = ResumeEncoder()
    print("Model created!")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")