app/main.py

from fastapi import FastAPI
from pydantic import BaseModel
import torch
import torch.nn.functional as F
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

from .model_def import BuildTransformer

app = FastAPI(title="Hindi-English Translator API")
model = None
tokenizer = None
device = torch.device("cpu")


from typing import Optional, List


class GreedyTranslationRequest(BaseModel):
    text: str


class BeamTranslationRequest(BaseModel):
    text: str
    beam_size: int = 3


class BeamAlternative(BaseModel):
    text: str
    score: float


class AttentionData(BaseModel):
    cross_attn: List[List[List[List[float]]]]  # [6 layers][8 heads][tgt_len][src_len]
    src_tokens: List[str]
    tgt_tokens: List[str]


class GreedyResponse(BaseModel):
    translated_text: str
    attention_weights: Optional[AttentionData] = None


class BeamResponse(BaseModel):
    translated_text: str
    alternatives: List[BeamAlternative]
    attention_weights: Optional[AttentionData] = None


@app.on_event("startup")
def load_assets():

    global model, tokenizer, device
    local_cache_dir = "/tmp/hf_cache"
    model_file = hf_hub_download(
        repo_id="Kush26/Transformer_Translation",
        filename="model.pth",
        cache_dir=local_cache_dir,
    )
    tokenizer_file = hf_hub_download(
        repo_id="Kush26/Transformer_Translation",
        filename="hindi-english_bpe_tokenizer.json",
        cache_dir=local_cache_dir,
    )
    tokenizer = Tokenizer.from_file(tokenizer_file)
    vocab_size = tokenizer.get_vocab_size()
    config = {
        "d_model": 256,
        "num_layers": 6,
        "num_heads": 8,
        "d_ff": 2048,
        "dropout": 0.1,
        "max_seq_len": 512,
    }
    model = BuildTransformer(
        src_vocab_size=vocab_size,
        trg_vocab_size=vocab_size,
        src_seq_len=config["max_seq_len"],
        trg_seq_len=config["max_seq_len"],
        d_model=config["d_model"],
        N=config["num_layers"],
        h=config["num_heads"],
        dropout=config["dropout"],
        d_ff=config["d_ff"],
    ).to(device)
    checkpoint = torch.load(model_file, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()


# --- Helpers ---
def get_all_token_strings(ids: list):
    """Return token strings for ALL positions (including special tokens) so labels match attention tensor dims."""
    tokens = []
    for tid in ids:
        tok_str = tokenizer.id_to_token(tid)
        if tok_str is not None:
            tokens.append(tok_str)
        else:
            tokens.append(f"<{tid}>")
    return tokens


def validate_request(text: str, beam_size: int = None):
    if not text or not text.strip():
        raise ValueError("Input text cannot be empty")
    if len(text) > 1000:
        raise ValueError("Input text exceeds maximum length of 1000 characters")
    if beam_size is not None:
        if beam_size < 1:
            raise ValueError("beam_size must be at least 1")
        if beam_size > 10:
            raise ValueError("beam_size cannot exceed 10")


# --- Translation Logic ---
def greedy_decode(sentence: str, max_len=100):
    PAD_token = tokenizer.token_to_id("[PAD]")
    model.eval()
    with torch.inference_mode():
        src_ids = (
            [tokenizer.token_to_id("[SOS]")]
            + tokenizer.encode(sentence).ids
            + [tokenizer.token_to_id("[EOS]")]
        )
        src_tensor = torch.tensor(src_ids).unsqueeze(0).to(device)
        src_mask = (src_tensor != PAD_token).unsqueeze(1).unsqueeze(2)
        encoder_output = model.encode(src_tensor, src_mask)
    tgt_tokens = [tokenizer.token_to_id("[SOS]")]

    all_cross_attn = []

    for _ in range(max_len):
        tgt_tensor = torch.tensor(tgt_tokens).unsqueeze(0).to(device)
        trg_mask_padding = (tgt_tensor != PAD_token).unsqueeze(1).unsqueeze(2)
        subsequent_mask = torch.tril(
            torch.ones(1, tgt_tensor.size(1), tgt_tensor.size(1), device=device)
        ).bool()
        trg_mask = trg_mask_padding & subsequent_mask
        with torch.inference_mode():
            decoder_output = model.decode(
                encoder_output, src_mask, tgt_tensor, trg_mask
            )
            logits = model.project(decoder_output)
        pred_token = logits.argmax(dim=-1)[0, -1].item()

        attn = model.get_cross_attention_weights()
        step_attn = []
        for layer_attn in attn:
            step_attn.append(layer_attn[:, -1, :].tolist()) 
        all_cross_attn.append(step_attn)

        tgt_tokens.append(pred_token)
        if pred_token == tokenizer.token_to_id("[EOS]"):
            break

    text = tokenizer.decode(tgt_tokens, skip_special_tokens=True)

    # Stack all step-level attention into a full [6][heads][tgt_len][src_len] matrix
    # all_cross_attn: [tgt_len] items, each [layers][heads][src_len]
    tgt_len = len(all_cross_attn)
    num_layers = 6
    if tgt_len > 0:
        src_len = len(all_cross_attn[0][0][0])
        cross_attn = []
        for l in range(num_layers):
            layer_mat = []
            for h in range(len(all_cross_attn[0][l])):
                head_mat = []
                for t in range(tgt_len):
                    head_mat.append(
                        [all_cross_attn[t][l][h][s] for s in range(src_len)]
                    )
                layer_mat.append(head_mat)
            cross_attn.append(layer_mat)
    else:
        cross_attn = []

    src_tokens = get_all_token_strings(src_ids)
    tgt_tokens_all = get_all_token_strings(tgt_tokens)

    return text, cross_attn, src_tokens, tgt_tokens_all


def beam_search_decode(sentence: str, beam_size: int, max_len=50):
    pad_token_id = tokenizer.token_to_id("[PAD]")
    model.eval()
    src_ids = (
        [tokenizer.token_to_id("[SOS]")]
        + tokenizer.encode(sentence).ids
        + [tokenizer.token_to_id("[EOS]")]
    )
    src_tensor = torch.tensor(src_ids).unsqueeze(0).to(device)
    src_mask = (src_tensor != pad_token_id).unsqueeze(1).unsqueeze(2)

    with torch.inference_mode():
        encoder_output = model.encode(src_tensor, src_mask)

    initial_beam = (
        torch.tensor([tokenizer.token_to_id("[SOS]")], device=device),
        0.0,
        1,
    )
    beams = [initial_beam]

    for _ in range(max_len):
        new_beams = []
        all_completed = True
        for seq, score, length in beams:
            if seq[-1].item() == tokenizer.token_to_id("[EOS]"):
                new_beams.append((seq, score, length))
                continue

            all_completed = False
            tgt_tensor = seq.unsqueeze(0)
            trg_mask_padding = (tgt_tensor != pad_token_id).unsqueeze(1).unsqueeze(2)
            subsequent_mask = torch.tril(
                torch.ones(1, tgt_tensor.size(1), tgt_tensor.size(1), device=device)
            ).bool()
            trg_mask = trg_mask_padding & subsequent_mask

            with torch.inference_mode():
                decoder_output = model.decode(
                    encoder_output, src_mask, tgt_tensor, trg_mask
                )
                logits = model.project(decoder_output)

            last_token_logits = logits[0, -1, :]
            log_probs = F.log_softmax(last_token_logits, dim=-1)
            top_log_probs, top_next_tokens = torch.topk(log_probs, beam_size)

            for i in range(beam_size):
                next_token = top_next_tokens[i]
                log_prob = top_log_probs[i].item()
                new_seq = torch.cat([seq, next_token.unsqueeze(0)])
                new_length = length + 1
                normalized_score = (score + log_prob) / new_length
                new_beams.append((new_seq, normalized_score, new_length))

        if all_completed:
            break

        new_beams.sort(key=lambda x: x[1], reverse=True)
        beams = new_beams[:beam_size]

    alternatives = []
    for seq, score, length in beams:
        decoded = tokenizer.decode(seq.tolist(), skip_special_tokens=True)
        alternatives.append({"text": decoded, "score": round(score, 4)})

    if not alternatives:
        return "", [], [], [], []

    best_text = alternatives[0]["text"]
    best_seq = beams[0][0]

    best_ids = best_seq.tolist()
    all_cross_attn = []
    attn_tgt = [best_ids[0]]
    for pos in range(1, len(best_ids)):
        attn_tensor = torch.tensor(attn_tgt).unsqueeze(0).to(device)
        attn_mask = (attn_tensor != pad_token_id).unsqueeze(1).unsqueeze(2)
        attn_subsequent = torch.tril(
            torch.ones(1, attn_tensor.size(1), attn_tensor.size(1), device=device)
        ).bool()
        with torch.inference_mode():
            model.decode(
                encoder_output, src_mask, attn_tensor, attn_mask & attn_subsequent
            )
        attn = model.get_cross_attention_weights()
        step_attn = []
        for layer_attn in attn:
            step_attn.append(layer_attn[:, -1, :].tolist()) 
        all_cross_attn.append(step_attn)
        attn_tgt.append(best_ids[pos])

    # Stack: [6][heads][tgt_len][src_len]
    tgt_len = len(all_cross_attn)
    num_layers = 6
    if tgt_len > 0:
        src_len = len(all_cross_attn[0][0][0])
        cross_attn = []
        for l in range(num_layers):
            layer_mat = []
            for h in range(len(all_cross_attn[0][l])):
                head_mat = []
                for t in range(tgt_len):
                    head_mat.append(
                        [all_cross_attn[t][l][h][s] for s in range(src_len)]
                    )
                layer_mat.append(head_mat)
            cross_attn.append(layer_mat)
    else:
        cross_attn = []

    src_tokens = get_all_token_strings(src_ids)
    tgt_tokens_all = get_all_token_strings(best_ids)
    return best_text, alternatives, cross_attn, src_tokens, tgt_tokens_all


@app.get("/")
def read_root():
    return {"message": "Hindi-English Translator API"}


@app.post("/translate/greedy", response_model=GreedyResponse)
def translate_greedy_endpoint(request: GreedyTranslationRequest):
    validate_request(request.text)
    translated_text, cross_attn, src_tokens, tgt_tokens = greedy_decode(request.text)
    return {
        "translated_text": translated_text,
        "attention_weights": {
            "cross_attn": cross_attn,
            "src_tokens": src_tokens,
            "tgt_tokens": tgt_tokens,
        },
    }


@app.post("/translate/beam", response_model=BeamResponse)
def translate_beam_endpoint(request: BeamTranslationRequest):
    validate_request(request.text, request.beam_size)
    translated_text, alternatives, cross_attn, src_tokens, tgt_tokens = (
        beam_search_decode(request.text, request.beam_size)
    )
    return {
        "translated_text": translated_text,
        "alternatives": alternatives,
        "attention_weights": {
            "cross_attn": cross_attn,
            "src_tokens": src_tokens,
            "tgt_tokens": tgt_tokens,
        },
    }