| """humanize-text-model inference wrapper (bilingual: English + Chinese). |
| |
| Loads the Lynote bilingual humanizer — two small seq2seq checkpoints in one |
| repo (``en/`` T5-small, ``zh/`` Chinese T5-small) — and routes each input to |
| the right language model. URLs, numbers, paths, code and quoted strings are |
| protected with tokenizer-friendly placeholders before generation and restored |
| afterwards. |
| |
| Usage: |
| from humanize import Humanizer |
| h = Humanizer() # loads "Lynote/humanize-text-model" from the Hub |
| print(h("It is important to note that this robust solution serves as a testament to our commitment.")) |
| """ |
|
|
| import argparse |
| import os |
| import re |
| import sys |
| from typing import List, Optional |
|
|
| import torch |
| from transformers import BertTokenizer, T5ForConditionalGeneration, T5Tokenizer |
|
|
| DEFAULT_MODEL = "Lynote/humanize-text-model" |
| MAX_CHARS = 4_000 |
| MAX_TOKENS = 256 |
| DEVICE = "mps" if torch.backends.mps.is_available() else ( |
| "cuda" if torch.cuda.is_available() else "cpu" |
| ) |
|
|
| PROTECTED_PATTERN = re.compile( |
| r"(`[^`]+`|https?://[^\s\u4e00-\u9fff，。！？；：、]+|(?:[A-Za-z]:)?(?:[/\\][\w.\-]+)+|" |
| r"\b\d+(?:\.\d+)?(?:%|[A-Za-z]+)?\b|\"[^\"]*\"|'[^']*')" |
| ) |
| CJK_RATIO = re.compile(r"[\u4e00-\u9fff]") |
|
|
|
|
| def protect(text: str, language: Optional[str] = None): |
| """Replace protected spans with language-aware placeholders: |
| EN ``PROTECTED_n``, ZH ``【保护n】`` (both tokenizer-friendly).""" |
| protected = [] |
| zh = language == "zh" or (language is None and is_chinese(text)) |
|
|
| def replace(match): |
| token = f"【保护{len(protected)}】" if zh else f"PROTECTED_{len(protected)}" |
| protected.append(match.group(0)) |
| return token |
|
|
| return PROTECTED_PATTERN.sub(replace, text), protected |
|
|
|
|
| def restore(text: str, protected: List[str]) -> str: |
| """Restore placeholders (either format, tolerant of inserted spaces).""" |
| for index, value in enumerate(protected): |
| text = re.sub(rf"PROTECTED_{index}\b", lambda m: value, text) |
| text = re.sub( |
| rf"【\s*保\s*护\s*{index}\s*】", lambda m: value, text |
| ) |
| return text |
|
|
|
|
| def is_chinese(text: str, threshold: float = 0.05) -> bool: |
| total = len(re.sub(r"\s+", "", text)) |
| if total == 0: |
| return False |
| return len(CJK_RATIO.findall(text)) / total > threshold |
|
|
|
|
| class Humanizer: |
| def __init__( |
| self, |
| model_id: str = DEFAULT_MODEL, |
| device: Optional[str] = None, |
| cache_dir: Optional[str] = None, |
| local_files_only: bool = False, |
| hf_token: Optional[str] = None, |
| lazy: bool = True, |
| ): |
| """model_id is a repo (or local dir) containing ``en/`` and ``zh/`` |
| subdirectories with separate checkpoints. With ``lazy=True`` (default) |
| a language model is only loaded when first needed.""" |
| self.model_id = model_id |
| self.device = torch.device(device or DEVICE) |
| self._models = {} |
| self._tokenizers = {} |
| self._cache_dir = cache_dir |
| self._local_only = local_files_only |
| self._token = hf_token |
| if not lazy: |
| self._get("en") |
| self._get("zh") |
|
|
| def _load(self, language: str, model: T5ForConditionalGeneration, tokenizer: T5Tokenizer): |
| self._models[language] = model |
| self._tokenizers[language] = tokenizer |
|
|
| def _get(self, language: str): |
| if language not in self._models: |
| model = T5ForConditionalGeneration.from_pretrained( |
| self.model_id, subfolder=language, |
| cache_dir=self._cache_dir, |
| local_files_only=self._local_only, token=self._token, |
| ) |
| tokenizer_cls = BertTokenizer if language == "zh" else T5Tokenizer |
| tokenizer = tokenizer_cls.from_pretrained( |
| self.model_id, subfolder=language, |
| cache_dir=self._cache_dir, |
| local_files_only=self._local_only, token=self._token, |
| ) |
| model.to(self.device).eval() |
| self._load(language, model, tokenizer) |
| return self._models[language], self._tokenizers[language] |
|
|
| def humanize(self, text: str, num_beams: int = 3, do_sample: bool = False) -> str: |
| text = (text or "").strip() |
| if not text: |
| raise ValueError("Input text is empty.") |
| if len(text) > MAX_CHARS: |
| raise ValueError(f"Input exceeds {MAX_CHARS:,} characters.") |
| return self.humanize_batch([text], num_beams=num_beams, do_sample=do_sample)[0] |
|
|
| def humanize_batch( |
| self, |
| texts: List[str], |
| num_beams: int = 3, |
| do_sample: bool = False, |
| batch_size: int = 8, |
| ) -> List[str]: |
| """Humanize a list of texts. URLs, numbers, paths, code and quotes are |
| protected per-text with placeholders and restored after generation.""" |
| masked = [] |
| protected_all = [] |
| for t in texts: |
| masked_text, protected = protect(t) |
| masked.append(masked_text) |
| protected_all.append(protected) |
| outputs: List[str] = [] |
| grouped: dict = {} |
| for i, t in enumerate(masked): |
| grouped.setdefault("zh" if is_chinese(t) else "en", []).append(i) |
| for language, indices in grouped.items(): |
| model, tokenizer = self._get(language) |
| for start in range(0, len(indices), batch_size): |
| chunk_idx = indices[start : start + batch_size] |
| enc = tokenizer( |
| [masked[i] for i in chunk_idx], |
| max_length=MAX_TOKENS, |
| truncation=True, |
| padding=True, |
| return_tensors="pt", |
| ) |
| if "token_type_ids" in enc: |
| enc.pop("token_type_ids") |
| enc = enc.to(self.device) |
| with torch.no_grad(): |
| gen_kwargs = dict( |
| num_beams=num_beams, |
| do_sample=do_sample, |
| early_stopping=True, |
| ) |
| if language == "zh": |
| |
| |
| |
| gen_kwargs["max_length"] = 40 |
| gen_kwargs["no_repeat_ngram_size"] = 4 |
| else: |
| gen_kwargs["max_length"] = MAX_TOKENS |
| gen = model.generate(**enc, **gen_kwargs) |
| decoded = tokenizer.batch_decode(gen, skip_special_tokens=True) |
| for j, text in zip(chunk_idx, decoded): |
| outputs.append((j, text)) |
| outputs.sort(key=lambda x: x[0]) |
|
|
| return [ |
| restore(self._polish(out), protected) |
| for out, protected in zip([o[1] for o in outputs], protected_all) |
| ] |
|
|
| @staticmethod |
| def _polish(text: str) -> str: |
| text = re.sub(r"^\s*(?:extra\d+\s*[,，。\s]*)+", "", text) |
| text = re.sub(r"(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff，。；：！？])", "", text) |
| text = re.sub(r"(?<=[，。；：！？])\s+(?=[\u4e00-\u9fff])", "", text) |
| text = re.sub(r"[ \t]{2,}", " ", text) |
| text = re.sub(r"\s+([,.;:!?，。；：！？])", r"\1", text) |
| text = re.sub(r"[,;]\s*\.", ".", text) |
| text = re.sub(r"\.\s*[,;]", ".", text) |
| text = re.sub(r"[，,]\s*(?=[。！？.!?])", "", text) |
| text = re.sub(r"\.\s+([a-z])", lambda m: ". " + m.group(1).upper(), text) |
| return text.strip(" ，") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Bilingual text humanizer (T5-small EN + Chinese T5-small)") |
| parser.add_argument("--text", help="text to humanize") |
| parser.add_argument("--input", help="path to a text file") |
| parser.add_argument("--output", help="write result to a file") |
| parser.add_argument("--model", default=DEFAULT_MODEL, help="model id or local dir") |
| parser.add_argument("--beams", type=int, default=3) |
| parser.add_argument("--list-models", action="store_true", help="print model ids and exit") |
| args = parser.parse_args() |
|
|
| if args.list_models: |
| print(DEFAULT_MODEL) |
| return |
|
|
| if args.text and args.input: |
| parser.error("provide only one of --text / --input") |
| source = args.text or (open(args.input, encoding="utf-8").read() if args.input else None) |
| if not source: |
| parser.error("provide --text or --input") |
|
|
| humanizer = Humanizer(model_id=args.model) |
| result = humanizer.humanize(source, num_beams=args.beams) |
| if args.output: |
| with open(args.output, "w", encoding="utf-8") as f: |
| f.write(result) |
| print(f"written to {args.output}") |
| else: |
| print(result) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
