"""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": # Small Chinese T5 degenerates on long sequences: cap # length and block exact 4-gram repetition (EN does not # need these). 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()