1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from transformers import AutoModel, AutoTokenizer
5from peft import get_peft_model, LoraConfig, TaskType
6from huggingface_hub import hf_hub_download
7
8class DhvaniV7(nn.Module):
9 def __init__(self, cfg):
10 super().__init__()
11 base = AutoModel.from_pretrained(
12 cfg['base_model'], torch_dtype=torch.bfloat16,
13 attn_implementation='eager', trust_remote_code=True
14 )
15 lora_config = LoraConfig(
16 r=cfg['lora_r'], lora_alpha=cfg['lora_alpha'],
17 lora_dropout=cfg['lora_dropout'],
18 target_modules=cfg['lora_targets'],
19 bias='none', task_type=TaskType.FEATURE_EXTRACTION
20 )
21 self.base = get_peft_model(base, lora_config)
22 self.trunk = nn.Sequential(
23 nn.Linear(cfg['hidden_dim'], cfg['trunk_dim']),
24 nn.LayerNorm(cfg['trunk_dim']), nn.GELU(),
25 )
26 self.surface_head = nn.Sequential(
27 nn.Linear(cfg['trunk_dim'], cfg['subspace_dim']),
28 nn.LayerNorm(cfg['subspace_dim']),
29 )
30 self.abhida_head = nn.Sequential(
31 nn.Linear(cfg['trunk_dim'], cfg['subspace_dim']),
32 nn.LayerNorm(cfg['subspace_dim']),
33 )
34 self.vyanjana_head = nn.Sequential(
35 nn.Linear(cfg['trunk_dim'], cfg['subspace_dim']),
36 nn.LayerNorm(cfg['subspace_dim']),
37 )
38
39 @staticmethod
40 def mean_pool(hidden, mask):
41 m = mask.unsqueeze(-1).float()
42 return (hidden * m).sum(1) / m.sum(1).clamp(min=1e-9)
43
44 def encode_tokens(self, input_ids, attention_mask):
45 out = self.base(input_ids=input_ids, attention_mask=attention_mask)
46 pooled = self.mean_pool(out.last_hidden_state.float(), attention_mask)
47 trunk = self.trunk(pooled)
48 return {
49 'surface': F.normalize(self.surface_head(trunk), p=2, dim=-1),
50 'abhida': F.normalize(self.abhida_head(trunk), p=2, dim=-1),
51 'vyanjana': F.normalize(self.vyanjana_head(trunk), p=2, dim=-1),
52 'full': F.normalize(torch.cat([
53 self.surface_head(trunk),
54 self.abhida_head(trunk),
55 self.vyanjana_head(trunk),
56 ], dim=-1), p=2, dim=-1),
57 }
58
59# Load
60ckpt_path = hf_hub_download(repo_id="rb512/dhvani-v7", filename="v7_best.pt")
61ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
62cfg = ckpt["config"]
63
64tokenizer = AutoTokenizer.from_pretrained(cfg['base_model'], trust_remote_code=True)
65if tokenizer.pad_token is None:
66 tokenizer.pad_token = tokenizer.eos_token
67
68model = DhvaniV7(cfg)
69model.base.load_state_dict(ckpt["lora"])
70model.trunk.load_state_dict(ckpt["trunk"])
71model.surface_head.load_state_dict(ckpt["surface_head"])
72model.abhida_head.load_state_dict(ckpt["abhida_head"])
73model.vyanjana_head.load_state_dict(ckpt["vyanjana_head"])
74model.eval()
75
76# Encode
77texts = ["The cat sat on the mat.", "A feline rested upon the rug."]
78enc = tokenizer(texts, max_length=128, truncation=True, padding='max_length', return_tensors='pt')
79with torch.no_grad():
80 embs = model.encode_tokens(enc['input_ids'], enc['attention_mask'])
81
82# Surface: high similarity (paraphrases)
83# Abhida: high similarity (same meaning, decorrelated from surface)
84# Vyanjana: similar (same register)
85print(f"Surface sim: {(embs['surface'][0] @ embs['surface'][1]).item():.3f}")
86print(f"Abhida sim: {(embs['abhida'][0] @ embs['abhida'][1]).item():.3f}")
87print(f"Vyanjana sim: {(embs['vyanjana'][0] @ embs['vyanjana'][1]).item():.3f}")