otelk/mt-error-detect-correct

MT Error Detection and Correction Models

This repository contains models for machine translation error detection and error correction.

Repository Structure

error_detection/
├── model/   # Model for translation error detection
└── lora/    # GRPO LoRA adapter for error detection

error_correction/
├── model/   # Model for translation error correction
└── lora/    # GRPO LoRA adapter for error correction

Usage

The repository contains two components:

• error_detection: machine translation error detection
• error_correction: machine translation error correction

Installation

pip install transformers huggingface_hub vllm

Load a Model

First download the repository:

import os
from huggingface_hub import snapshot_download
from transformers import AutoTokenizer
from vllm import LLM
from vllm.lora.request import LoRARequest

repo_dir = snapshot_download(
    repo_id="otelk/mt-error-detect-correct"
)

For error detection:

model_dir = os.path.join(repo_dir, "error_detection", "model")
adapter_dir = os.path.join(repo_dir, "error_detection", "lora")

For error correction:

model_dir = os.path.join(repo_dir, "error_correction", "model")
adapter_dir = os.path.join(repo_dir, "error_correction", "lora")

tokenizer = AutoTokenizer.from_pretrained(
    model_dir,
    trust_remote_code=True
)

llm = LLM(
    model=model_dir,
    tokenizer=model_dir,
    dtype="bfloat16",
    enable_lora=True,
    max_lora_rank=16,
    trust_remote_code=True
)

lora_request = LoRARequest(
    "mt_lora",
    1,
    adapter_dir
)

The model can then be used with vLLM by passing the corresponding lora_request during generation.

outputs = llm.generate(
    prompts,
    sampling_params=sampling_params,
    lora_request=lora_request
)