app.py

# import torch
# import torch.nn as nn
# import torch.nn.functional as F
# import gradio as gr
# from torch.utils.data import Dataset
# from transformers import PreTrainedModel, PretrainedConfig, Trainer, TrainingArguments
# from datasets import load_dataset
# import numpy as np

# # =====================
# # 1. Load Dataset Subsets
# # =====================
# dataset = load_dataset("bashyaldhiraj2067/500k_copy_error_dataset")
# train_subset = dataset["train"].select(range(int(len(dataset["train"]) * 0.1)))
# test_subset = dataset["test"].select(range(int(len(dataset["test"]) * 0.1)))
# print(f"Subset train size: {len(train_subset)}")
# print(f"Subset test size: {len(test_subset)}")

# # =====================
# # 2. Tokenizer
# # =====================
# special_tokens = ["<pad>", "<s>", "</s>", "<unk>"]
# nepali_chars = list("अआइईउऊऋॠऌॡऎएऐओऔकखगघङचछजझञटठडढणतथदधनपफबभमयरलवशषसह्ािीुूृॄेैोौंंःँ।०१२३४५६७८९,.;?!़ॅंःॊॅऒऽॉड़ॐ॥ऑऱफ़ढ़")
# char_vocab = special_tokens + nepali_chars
# char2id = {char: idx for idx, char in enumerate(char_vocab)}
# id2char = {idx: char for char, idx in char2id.items()}
# vocab_size = len(char2id)

# class CharTokenizer:
#     def __init__(self, char2id, id2char, vocab_size):
#         self.char2id = char2id
#         self.id2char = id2char
#         self.pad_token_id = char2id["<pad>"]
#         self.unk_token_id = char2id["<unk>"]
#         self.bos_token_id = char2id["<s>"]
#         self.eos_token_id = char2id["</s>"]
#         self.vocab_size = vocab_size

#     def encode(self, text, max_length=128):
#         ids = [self.char2id.get(ch, self.unk_token_id) for ch in text]
#         ids = ids[:max_length]
#         return ids + [self.pad_token_id] * (max_length - len(ids))

#     def decode(self, ids):
#         return ''.join([self.id2char.get(i, '') for i in ids if i != self.pad_token_id])

#     def __call__(self, text, text_target=None, max_length=128):
#         input_ids = self.encode(text, max_length)
#         input_ids = torch.clamp(torch.tensor(input_ids), max=self.vocab_size - 1).tolist()
#         result = {"input_ids": input_ids, "attention_mask": [1 if i != self.pad_token_id else 0 for i in input_ids]}
#         if text_target:
#             labels = self.encode(text_target, max_length)
#             result["labels"] = labels
#         return result

# tokenizer = CharTokenizer(char2id, id2char, vocab_size=vocab_size)

# # =====================
# # 3. Dataset
# # =====================
# class CopyDataset(Dataset):
#     def __init__(self, data, tokenizer, max_length=128):
#         self.data = data
#         self.tokenizer = tokenizer
#         self.max_length = max_length

#     def __len__(self):
#         return len(self.data)

#     def __getitem__(self, idx):
#         noisy = self.data[idx]['incorrect']
#         clean = self.data[idx]['correct']
#         return self.tokenizer(noisy, text_target=clean, max_length=self.max_length)

# train_dataset = CopyDataset(train_subset, tokenizer)
# eval_dataset = CopyDataset(test_subset, tokenizer)

# # =====================
# # 4. Transformer with Copy Mechanism
# # =====================
# class TransformerCopyConfig(PretrainedConfig):
#     def __init__(self, vocab_size=len(char2id), **kwargs):
#         super().__init__(**kwargs)
#         self.vocab_size = vocab_size

# # --- Model Components ---
# class PositionalEncoding(nn.Module):
#     def __init__(self, d_model, max_len=512):
#         super().__init__()
#         pe = torch.zeros(max_len, d_model)
#         position = torch.arange(0, max_len).unsqueeze(1)
#         div_term = torch.exp(torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model))
#         pe[:, 0::2] = torch.sin(position * div_term)
#         pe[:, 1::2] = torch.cos(position * div_term)
#         self.register_buffer('pe', pe.unsqueeze(0))

#     def forward(self, x):
#         return x + self.pe[:, :x.size(1)]

# class TransformerCopyModel(nn.Module):
#     def __init__(self, vocab_size, d_model=256, nhead=8, num_layers=4, dim_ff=512, dropout=0.1):
#         super().__init__()
#         self.embedding = nn.Embedding(vocab_size, d_model)
#         self.positional_encoding = PositionalEncoding(d_model)

#         encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_ff, dropout)
#         decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_ff, dropout)

#         self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
#         self.decoder = nn.TransformerDecoder(decoder_layer, num_layers)

#         self.copy_attention = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
#         self.copy_gate = nn.Linear(d_model * 2, 1)

#         self.output_layer = nn.Linear(d_model, vocab_size)

#     def forward(self, input_ids, attention_mask=None, labels=None):
#         src = input_ids
#         tgt = labels[:, :-1]
#         tgt_y = labels[:, 1:]

#         src_embed = self.embedding(src)
#         tgt_embed = self.embedding(tgt)
#         src_embed = self.positional_encoding(src_embed)
#         tgt_embed = self.positional_encoding(tgt_embed)

#         src_mask = (src == tokenizer.pad_token_id)
#         tgt_mask = (tgt == tokenizer.pad_token_id)

#         memory = self.encoder(src_embed.transpose(0, 1), src_key_padding_mask=src_mask)
#         output = self.decoder(
#             tgt_embed.transpose(0, 1),
#             memory,
#             tgt_key_padding_mask=tgt_mask,
#             memory_key_padding_mask=src_mask
#         )

#         attn_output, attn_weights = self.copy_attention(output, memory, memory, key_padding_mask=src_mask)
#         concat = torch.cat([output, attn_output], dim=-1)
#         copy_prob = torch.sigmoid(self.copy_gate(concat))

#         gen_logits = self.output_layer(output)
#         gen_probs = F.softmax(gen_logits, dim=-1)

#         loss = F.cross_entropy(
#             gen_logits.transpose(0, 1).reshape(-1, gen_logits.size(-1)),
#             tgt_y.reshape(-1),
#             ignore_index=tokenizer.pad_token_id
#         ) if labels is not None else None

#         return {"loss": loss, "logits": gen_logits.transpose(0, 1)}

# # --- HF Wrapper ---
# class TransformerCopyHF(PreTrainedModel):
#     config_class = TransformerCopyConfig
#     def __init__(self, config):
#         super().__init__(config)
#         self.model = TransformerCopyModel(config.vocab_size)

#     def forward(self, input_ids, attention_mask=None, labels=None):
#         return self.model(input_ids, attention_mask, labels)

# model = TransformerCopyHF.from_pretrained("bashyaldhiraj2067/remove_copy_transformer")
# model.eval()

# # =====================
# # 5. Inference Function
# # =====================
# def generate_clean_text(input_text, max_length=128):
#     model_input = tokenizer.encode(input_text, max_length=max_length)
#     input_ids = torch.tensor([model_input])
#     # Create dummy target input (just start token)
#     decoder_input = torch.tensor([[tokenizer.bos_token_id]])
#     output_tokens = []
#     for _ in range(max_length):
#         with torch.no_grad():
#             out = model(input_ids=input_ids, labels=torch.cat([decoder_input, torch.zeros((1, 1), dtype=torch.long)], dim=1))
#             next_token_logits = out["logits"][:, -1, :]
#             next_token = torch.argmax(next_token_logits, dim=-1)

#         next_token_id = next_token.item()
       
#         if next_token_id == tokenizer.pad_token_id:
#             break
#         output_tokens.append(next_token_id)
#         decoder_input = torch.cat([decoder_input, next_token.unsqueeze(0)], dim=1)

#     return tokenizer.decode(output_tokens)


# # Gradio Interface Setup
# iface = gr.Interface(
#     fn=generate_clean_text,
#     inputs=gr.Textbox(label="Noisy Text"),
#     outputs=gr.Textbox(label="Cleaned Text"),
#     live=True
# )

# iface.launch(debug=True)
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from transformers import PreTrainedModel, PretrainedConfig

# =========================================================
# 1. Tokenizer (CUSTOM – REQUIRED)
# =========================================================
special_tokens = ["<pad>", "<s>", "</s>", "<unk>"]
nepali_chars = list(
    "अआइईउऊऋॠऌॡऎएऐओऔकखगघङचछजझञटठडढणतथदधनपफबभमयरलवशषसह"
    "ािीुूृॄेैोौंंःँ।०१२३४५६७८९,.;?!़ॅॊऒऽॉड़ॐ॥ऑऱफ़ढ़"
)
char_vocab = special_tokens + nepali_chars
char2id = {c: i for i, c in enumerate(char_vocab)}
id2char = {i: c for c, i in char2id.items()}

class CharTokenizer:
    def __init__(self):
        self.pad_token_id = char2id["<pad>"]
        self.unk_token_id = char2id["<unk>"]
        self.bos_token_id = char2id["<s>"]
        self.eos_token_id = char2id["</s>"]
        self.vocab_size = len(char2id)

    def encode(self, text, max_length=128):
        ids = [char2id.get(ch, self.unk_token_id) for ch in text]
        ids = ids[:max_length]
        return ids + [self.pad_token_id] * (max_length - len(ids))

    def decode(self, ids):
        return "".join(id2char.get(i, "") for i in ids if i != self.pad_token_id)


tokenizer = CharTokenizer()

# =========================================================
# 2. Model Definition (CUSTOM – REQUIRED)
# =========================================================
class TransformerCopyHF(PreTrainedModel):
    config_class = TransformerCopyConfig
    def __init__(self, config):
        super().__init__(config)
        self.model = TransformerCopyModel(
            vocab_size=config.vocab_size,
            d_model=256,
            nhead=8,
            num_layers=4,
            dim_ff=512
        )


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div = torch.exp(
            torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div)
        pe[:, 1::2] = torch.cos(position * div)
        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        return x + self.pe[:, : x.size(1)]

class TransformerCopyModel(nn.Module):
    def __init__(self, vocab_size, d_model=256, nhead=8, num_layers=4, dim_ff=512, dropout=0.1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos = PositionalEncoding(d_model)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_ff,
            dropout=dropout
        )
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model,
            nhead=nhead,
            dim_feedforward=dim_ff,
            dropout=dropout
        )

        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_layers)
        self.fc = nn.Linear(d_model, vocab_size)



    def forward(self, src, tgt):
        src_emb = self.pos(self.embedding(src))
        tgt_emb = self.pos(self.embedding(tgt))

        memory = self.encoder(src_emb.transpose(0, 1))
        out = self.decoder(
            tgt_emb.transpose(0, 1), memory
        )

        return self.fc(out.transpose(0, 1))

class TransformerCopyHF(PreTrainedModel):
    config_class = TransformerCopyConfig

    def __init__(self, config):
        super().__init__(config)
        self.model = TransformerCopyModel(config.vocab_size)

    def forward(self, input_ids, decoder_input_ids):
        return self.model(input_ids, decoder_input_ids)

# =========================================================
# 3. Load Weights from Hugging Face
# =========================================================
device = "cuda" if torch.cuda.is_available() else "cpu"

model = TransformerCopyHF.from_pretrained(
    "bashyaldhiraj2067/epoch15_nepali-bart-copy-mechanism"
).to(device)
model.eval()

# =========================================================
# 4. Inference Function
# =========================================================
@torch.no_grad()
def generate_clean_text(text, max_len=128):
    src = torch.tensor(
        [tokenizer.encode(text, max_len)],
        device=device
    )

    tgt = torch.tensor(
        [[tokenizer.bos_token_id]],
        device=device
    )

    output_tokens = []

    for _ in range(max_len):
        logits = model(src, tgt)
        next_token = torch.argmax(logits[:, -1], dim=-1)

        token_id = next_token.item()
        if token_id == tokenizer.pad_token_id:
            break

        output_tokens.append(token_id)
        tgt = torch.cat([tgt, next_token.unsqueeze(0)], dim=1)

    return tokenizer.decode(output_tokens)

# =========================================================
# 5. Gradio UI
# =========================================================
with gr.Blocks(title="Nepali GEC – Copy Transformer") as demo:
    gr.Markdown("## 🇳🇵 Nepali Grammatical Error Correction")

    inp = gr.Textbox(
        label="Noisy / Incorrect Text",
        lines=4,
        placeholder="यहाँ गलत नेपाली वाक्य लेख्नुहोस्"
    )

    out = gr.Textbox(
        label="Corrected Text",
        lines=4
    )

    btn = gr.Button("Correct")

    btn.click(
        fn=generate_clean_text,
        inputs=inp,
        outputs=out
    )

demo.launch()