Upload 6 files

README.md

@@ -1,14 +1 @@
----
-title: Machine Translation
-emoji: 🐠
-colorFrom: green
-colorTo: red
-sdk: gradio
-sdk_version: 5.43.1
-app_file: app.py
-pinned: false
-license: cc0-1.0
-short_description: translate english to vietnamese
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Machine_Translation

app.py

@@ -0,0 +1,71 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from io import BytesIO
+import requests
+from model import TransformerSeq2Seq,translate
+from utils import load_tokenizers_and_embeddings
+
+import torch
+
+   # class mô hình của bạn
+
+app = FastAPI()
+
+# ===== 1. Load model và tokenizer khi khởi động server =====
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# ===== Load 1 lần khi start server =====
+resources = load_tokenizers_and_embeddings()
+tokenizer_vi = resources["tokenizer_vi"]
+embedding_matrix_vi = resources["embedding_vi"]
+tokenizer_en = resources["tokenizer_en"]
+embedding_matrix_en = resources["embedding_en"]
+device = resources["device"]
+
+print("✅ Tokenizers & embeddings loaded!")
+if isinstance(embedding_matrix_en, torch.Tensor):
+    embed_dim = embedding_matrix_en.size(1)
+else:  # nn.Embedding
+    embed_dim = embedding_matrix_en.embedding_dim
+max_len = 128
+batch_size = 32
+# Load model
+model = TransformerSeq2Seq(
+    embed_dim=embed_dim,
+    vocab_size=tokenizer_vi.vocab_size,   # hoặc len(tokenizer_vi)
+    embedding_decoder=embedding_matrix_vi,      # embedding target đã có sẵn
+    num_heads=4,
+    num_layers=2,
+    dim_feedforward=256,
+    dropout=0.1,
+    freeze_decoder_emb=True,
+    max_len=max_len
+)
+MODEL_URL = "https://huggingface.co/nemabruh404/Machine_Translation/resolve/main/model_state_dict.pt"
+
+# Fetch model từ Hub
+checkpoint_bytes = BytesIO(requests.get(MODEL_URL).content)
+checkpoint = torch.load(checkpoint_bytes, map_location=device)
+
+# Load state dict
+model.load_state_dict(checkpoint["model_state_dict"])
+model.to(device)
+model.eval()
+
+print("✅ Model loaded from Hugging Face Hub")
+print("Model loaded")
+class TranslationRequest(BaseModel):
+    text: str
+# ===== Endpoint dịch =====
+@app.post("/translate")
+def translate_api(req: TranslationRequest):
+    output = translate(
+        model=model,
+        src_sentence=req.text,
+        tokenizer_src=tokenizer_en,   # tiếng Anh -> input
+        tokenizer_tgt=tokenizer_vi,   # tiếng Việt -> output
+        embedding_src=embedding_matrix_en,
+        device=device,
+        max_len=max_len
+    )
+    return {"input": req.text, "translation": output}

main.py

@@ -0,0 +1,7 @@
+from app import app
+import uvicorn
+import os
+
+port = int(os.environ.get("PORT", 10000))  # Render sẽ set PORT
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=port)

model.py

@@ -0,0 +1,151 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+import math
+
+# ---------------- Positional Encoding ----------------
+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).float()
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)  # (1, max_len, d_model)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        # x: (B, T, D)
+        return x + self.pe[:, :x.size(1)].to(x.device)
+
+# ---------------- Transformer (sửa để match training) ----------------
+class TransformerSeq2Seq(nn.Module):
+    """
+    Thiết kế sao cho forward(src_embedded, tgt_input_ids, src_attn_mask=None, tgt_attn_mask=None)
+    - src_embedded: (B, S, E) — bạn có thể pass embedding matrix bên ngoài (embedding_src[src_ids])
+    - tgt_input_ids: (B, T)  — token ids cho decoder input (BOS.. token_{n-1})
+    - src_attn_mask / tgt_attn_mask: (B, S) / (B, T) with 1 for real tokens, 0 for pad
+    """
+    def __init__(self,
+                 embed_dim,
+                 vocab_size,                   # target vocab size (output dim)
+                 embedding_decoder=None,       # pretrained weights (np array or torch.Tensor) or None
+                 num_heads=2,
+                 num_layers=2,
+                 dim_feedforward=256,
+                 dropout=0.1,
+                 freeze_decoder_emb=True,
+                 max_len=512):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.vocab_size = vocab_size
+
+        # positional encoding
+        self.pos_encoder = PositionalEncoding(embed_dim, max_len=max_len)
+
+        # decoder embedding (pretrained optional)
+        if embedding_decoder is None:
+            self.embedding_decoder = nn.Embedding(vocab_size, embed_dim)
+        else:
+            if not isinstance(embedding_decoder, torch.Tensor):
+                embedding_decoder = torch.tensor(embedding_decoder, dtype=torch.float)
+            self.embedding_decoder = nn.Embedding.from_pretrained(embedding_decoder, freeze=freeze_decoder_emb)
+
+        # encoder/decoder (batch_first True -> inputs shape (B, T, E))
+        self.encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads,
+                                                        dim_feedforward=dim_feedforward, dropout=dropout,
+                                                        batch_first=True)
+        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)
+
+        self.decoder_layer = nn.TransformerDecoderLayer(d_model=embed_dim, nhead=num_heads,
+                                                        dim_feedforward=dim_feedforward, dropout=dropout,
+                                                        batch_first=True)
+        self.decoder = nn.TransformerDecoder(self.decoder_layer, num_layers=num_layers)
+
+        self.output_proj = nn.Linear(embed_dim, vocab_size)
+
+    def forward(self, src_embedded, tgt_input_ids, src_attn_mask=None, tgt_attn_mask=None):
+        """
+        src_embedded : (B, S, E)
+        tgt_input_ids: (B, T)
+        src_attn_mask : (B, S) mask: 1 real token, 0 pad  (optional)
+        tgt_attn_mask : (B, T) same
+        """
+        device = src_embedded.device
+        # tgt embedding
+        tgt_embedded = self.embedding_decoder(tgt_input_ids)  # (B, T, E)
+
+        # add positional encoding
+        src = self.pos_encoder(src_embedded)  # (B, S, E)
+        tgt = self.pos_encoder(tgt_embedded)  # (B, T, E)
+
+        # prepare key_padding_mask: True at positions that should be masked (pad positions)
+        src_key_padding_mask = None
+        tgt_key_padding_mask = None
+        if src_attn_mask is not None:
+            src_key_padding_mask = (src_attn_mask == 0).to(device)  # (B, S), bool
+        if tgt_attn_mask is not None:
+            tgt_key_padding_mask = (tgt_attn_mask == 0).to(device)  # (B, T)
+
+        # encode
+        memory = self.encoder(src, src_key_padding_mask=src_key_padding_mask)  # (B, S, E)
+
+        # causal mask for decoder (T x T)
+        T = tgt.size(1)
+        if T > 0:
+            tgt_mask = torch.triu(torch.full((T, T), float('-inf'), device=device), diagonal=1)
+        else:
+            tgt_mask = None
+
+        # decode
+        output = self.decoder(tgt, memory,
+                              tgt_mask=tgt_mask,
+                              tgt_key_padding_mask=tgt_key_padding_mask,
+                              memory_key_padding_mask=src_key_padding_mask)  # (B, T, E)
+
+        logits = self.output_proj(output)  # (B, T, vocab)
+        return logits
+
+# ---------------- Helpers to apply embedding_src (tensor or nn.Embedding) ----------------
+def apply_src_embedding(embedding_src, src_ids):
+    """
+    embedding_src can be:
+      - torch.Tensor of shape (vocab_src, embed_dim)  -> indexing
+      - nn.Embedding instance -> call( ids )
+    src_ids: LongTensor (B, S)
+    return: (B, S, E) float tensor on same device as src_ids
+    """
+    if isinstance(embedding_src, nn.Embedding):
+        return embedding_src(src_ids)
+    else:
+        # assume it's a tensor/ndarray
+        if not isinstance(embedding_src, torch.Tensor):
+            embedding_src = torch.tensor(embedding_src, dtype=torch.float, device=src_ids.device)
+        else:
+            embedding_src = embedding_src.to(src_ids.device)
+        return embedding_src[src_ids]
+@torch.no_grad()
+def translate(model, src_sentence, tokenizer_src, tokenizer_tgt, embedding_src, device, max_len=50):
+    model.eval()
+    inputs = tokenizer_src([src_sentence], return_tensors="pt", padding=True, truncation=True, max_length=128)
+    src_ids = inputs["input_ids"].to(device)      # (1, S)
+    src_attn = inputs.get("attention_mask", None)
+    if src_attn is not None:
+        src_attn = src_attn.to(device)
+
+    src_embedded = apply_src_embedding(embedding_src, src_ids)  # (1, S, E)
+
+    decoded_ids = [tokenizer_tgt.cls_token_id]
+    for _ in range(max_len):
+        decoder_input = torch.tensor([decoded_ids], device=device)
+        # for decode we don't need tgt_attn_mask (we build causal mask inside model)
+        logits = model(src_embedded, decoder_input, src_attn_mask=src_attn, tgt_attn_mask=None)
+        next_token = logits[:, -1, :].argmax(dim=-1).item()
+        decoded_ids.append(next_token)
+        if next_token == tokenizer_tgt.sep_token_id:
+            break
+
+    return tokenizer_tgt.decode(decoded_ids, skip_special_tokens=True)

requirements.txt

@@ -0,0 +1,6 @@
+fastapi
+uvicorn[standard]
+torch
+transformers
+requests
+pydantic

utils.py

@@ -0,0 +1,23 @@
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def load_tokenizers_and_embeddings():
+    # ===== Vietnamese PhoBERT =====
+    tokenizer_vi = AutoTokenizer.from_pretrained("vinai/phobert-base")
+    model_vi = AutoModel.from_pretrained("vinai/phobert-base").to(device)
+    embedding_matrix_vi = model_vi.embeddings.word_embeddings.weight
+
+    # ===== English BERT =====
+    tokenizer_en = AutoTokenizer.from_pretrained("bert-base-cased-finetuned-mrpc")
+    model_en = AutoModel.from_pretrained("bert-base-cased-finetuned-mrpc").to(device)
+    embedding_matrix_en = model_en.embeddings.word_embeddings.weight
+    
+    return {
+        "tokenizer_vi": tokenizer_vi,
+        "embedding_vi": embedding_matrix_vi,
+        "tokenizer_en": tokenizer_en,
+        "embedding_en": embedding_matrix_en,
+        "device": device
+    }