| import os |
| import json |
| import time |
| from typing import Any, Dict, Optional, Union |
| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel, AutoModelForCausalLM, AutoTokenizer |
| from safetensors.torch import load_file |
| from peft import PeftModel |
| from .configuration_kutup import KutupConfig |
|
|
| class EvidenceRoutingBlock(nn.Module): |
| def __init__(self, d_model: int, n_heads: int = 8, dropout: float = 0.05): |
| super().__init__() |
| self.cross_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) |
| self.norm = nn.LayerNorm(d_model) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, query: torch.Tensor, key_value: torch.Tensor) -> torch.Tensor: |
| attn_out, _ = self.cross_attn(query, key_value, key_value) |
| return self.norm(query + self.dropout(attn_out)) |
|
|
| class JointSchemaHead(nn.Module): |
| def __init__(self, hidden_size: int = 2048, decision_width: int = 1024, num_layers: int = 4): |
| super().__init__() |
| self.proj = nn.Linear(hidden_size, decision_width) |
| self.routing = EvidenceRoutingBlock(decision_width) |
| encoder_layer = nn.TransformerEncoderLayer( |
| d_model=decision_width, |
| nhead=8, |
| dim_feedforward=decision_width * 4, |
| dropout=0.05, |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) |
| self.scorer = nn.Sequential( |
| nn.Linear(decision_width, decision_width // 2), |
| nn.GELU(), |
| nn.Linear(decision_width // 2, 1), |
| ) |
|
|
| def forward(self, prefix_emb: torch.Tensor, option_embs: torch.Tensor) -> torch.Tensor: |
| p_proj = self.proj(prefix_emb) |
| o_proj = self.proj(option_embs) |
| routed = self.routing(o_proj, p_proj) |
| joint = self.transformer(routed) |
| scores = self.scorer(joint).squeeze(-1) |
| return scores |
|
|
| class KutupForDecision(PreTrainedModel): |
| config_class = KutupConfig |
|
|
| def __init__(self, config: KutupConfig): |
| super().__init__(config) |
| self.config = config |
| self.head = JointSchemaHead( |
| hidden_size=config.hidden_size, |
| decision_width=config.decision_width |
| ) |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): |
| config = kwargs.pop("config", None) |
| if config is None: |
| config = KutupConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) |
| |
| device_map = kwargs.get("device_map", None) |
| torch_dtype = getattr(torch, config.torch_dtype) if isinstance(config.torch_dtype, str) else torch.bfloat16 |
| |
| model = cls(config) |
| |
| |
| print(f"[KUTUP-1] Backbone yukleniyor: {config.backbone_model_id}...") |
| backbone = AutoModelForCausalLM.from_pretrained( |
| config.backbone_model_id, |
| torch_dtype=torch_dtype, |
| device_map=device_map, |
| ) |
| |
| |
| lora_dir = os.path.join(pretrained_model_name_or_path, "lora_adapter") |
| if os.path.exists(lora_dir): |
| print(f"[KUTUP-1] LoRA yukleniyor: {lora_dir}...") |
| backbone = PeftModel.from_pretrained(backbone, lora_dir) |
| backbone = backbone.merge_and_unload() |
| model.backbone = backbone |
| |
| |
| head_path = os.path.join(pretrained_model_name_or_path, "joint_head.safetensors") |
| if os.path.exists(head_path): |
| print(f"[KUTUP-1] Joint Head yukleniyor: {head_path}...") |
| weights = load_file(head_path) |
| model.head.load_state_dict(weights, strict=False) |
| |
| model.to(torch_dtype) |
| model.eval() |
| return model |
|
|
| @torch.inference_mode() |
| def decide(self, tokenizer: AutoTokenizer, state: str, questions: Dict[str, Any]) -> Dict[str, Any]: |
| device = next(self.parameters()).device |
| answers = {} |
| for qid, qdef in questions.items(): |
| qtype = qdef.get("type", "choice") |
| inst = qdef.get("instructions", "") |
| crit = qdef.get("criteria", {}) |
| |
| prefix_text = f"[DURUM] |
| {state} |
| |
| [SORU] |
| {inst}" |
| p_inputs = tokenizer(prefix_text, return_tensors="pt", truncation=True, max_length=self.config.max_length).to(device) |
| p_out = self.backbone.model(**p_inputs) |
| p_emb = p_out.last_hidden_state |
| |
| opts = [] |
| if qtype == "noul": |
| opts = ["Evet", "Hayır"] |
| opt_keys = ["true", "false"] |
| elif isinstance(crit, dict): |
| opt_keys = list(crit.keys()) |
| opts = [f"{k}: {v}" for k, v in crit.items()] |
| elif isinstance(crit, list): |
| opt_keys = [str(i) for i in range(len(crit))] |
| opts = [f"Seviye {i}: {v}" for i, v in enumerate(crit)] |
| |
| opt_embs = [] |
| for opt in opts: |
| o_inp = tokenizer(f"[SECENEK] {opt}", return_tensors="pt").to(device) |
| o_out = self.backbone.model(**o_inp) |
| opt_embs.append(o_out.last_hidden_state.mean(dim=1)) |
| |
| all_opt_emb = torch.stack(opt_embs, dim=1) |
| logits = self.head(p_emb, all_opt_emb).squeeze(0) / self.config.temperature |
| probs = torch.softmax(logits, dim=-1).tolist() |
| |
| prob_dict = {k: round(p, 4) for k, p in zip(opt_keys, probs)} |
| best_idx = int(torch.argmax(logits).item()) |
| best_key = opt_keys[best_idx] |
| conf = round(probs[best_idx], 4) |
| |
| if qtype == "noul": |
| answers[qid] = {"type": "noul", "noul": prob_dict.get("true", 0.5), "confidence": conf} |
| else: |
| answers[qid] = {"choice": best_key, "confidence": conf, "probabilities": prob_dict} |
| |
| return {"model": "TurkishCodeMan/KUTUP-1-Decision-2B", "answers": answers} |
|
|
