Spaces:
Running on Zero
Running on Zero
| import spaces | |
| import json | |
| import os | |
| import re | |
| from threading import Thread | |
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForImageTextToText, AutoTokenizer, TextIteratorStreamer | |
| MODEL_ID = "abacusai/Smaug-Humanizer" | |
| MAX_WORDS = 2000 | |
| CHUNK_WORDS = 500 | |
| STOP = "<end_of_turn>" | |
| PROMPT = """Rewrite the AI-written text below so that it reads as if a person wrote it, in a person's own voice, wording and rhythm. | |
| Keep every fact, number, name and quote. Keep the structure and the language of the text. Output only the rewritten text. | |
| Text: | |
| {text}""" | |
| token = os.environ.get("HF_TOKEN") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=token) | |
| tokenizer.padding_side = "left" | |
| model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, dtype=torch.bfloat16, token=token).to("cuda") | |
| model.eval() | |
| def split_chunks(text: str, chunk_words: int = CHUNK_WORDS) -> list[str]: | |
| # ~chunk_words-word chunks at paragraph breaks; a very long paragraph is cut at sentence ends | |
| paras = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()] | |
| out, cur, n = [], [], 0 | |
| for p in paras: | |
| w = len(p.split()) | |
| if w > chunk_words * 1.5: | |
| if cur: | |
| out.append("\n\n".join(cur)) | |
| words, pieces, start = re.findall(r"\S+\s*", p), [[]], 0 | |
| for end, word in enumerate(words, 1): | |
| if end < len(words) and word.rstrip()[-1] not in ".!?": | |
| continue | |
| for i in range(start, end, chunk_words): | |
| part = words[i:min(i + chunk_words, end)] | |
| if pieces[-1] and len(pieces[-1]) + len(part) > chunk_words: | |
| pieces.append([]) | |
| pieces[-1] += part | |
| start = end | |
| if len(pieces) > 1 and len(pieces[-1]) < chunk_words * 0.4: | |
| tail = pieces.pop() | |
| pieces[-1] += tail | |
| out += ["".join(piece).strip() for piece in pieces[:-1]] | |
| cur, n = ["".join(pieces[-1]).strip()], len(pieces[-1]) | |
| continue | |
| if n + w > chunk_words and cur: | |
| out.append("\n\n".join(cur)) | |
| cur, n = [], 0 | |
| cur.append(p) | |
| n += w | |
| if cur and out and n < chunk_words * 0.4: | |
| out[-1] += "\n\n" + "\n\n".join(cur) | |
| elif cur: | |
| out.append("\n\n".join(cur)) | |
| return out | |
| def _prompt(chunk: str) -> str: | |
| return tokenizer.apply_chat_template([{"role": "user", "content": PROMPT.format(text=chunk)}], | |
| add_generation_prompt=True, tokenize=False) | |
| def _sampling(chunks: list[str]) -> dict: | |
| return dict(max_new_tokens=max(len(c.split()) for c in chunks) * 3 + 300, do_sample=True, temperature=1.0, top_p=0.95, | |
| top_k=64, stop_strings=[STOP], tokenizer=tokenizer) | |
| def _duration(text: str) -> int: | |
| chunks = split_chunks((text or "").strip()) or [""] | |
| return min(120, 30 + int(max(len(c.split()) for c in chunks) * 0.15) + 5 * len(chunks)) | |
| def humanize(text: str): | |
| """Rewrite AI-generated text so that it reads as if a person wrote it, keeping every fact, number and name. | |
| Texts over 500 words are split at paragraph breaks into chunks of about 500 words; the chunks are rewritten | |
| together and joined with a blank line. | |
| Args: | |
| text: The AI-written text to rewrite, up to about 2,000 words. | |
| """ | |
| text = (text or "").strip() | |
| if not text: | |
| raise gr.Error("Paste some text first.") | |
| if len(text.split()) > MAX_WORDS: | |
| raise gr.Error(f"Please keep the input under {MAX_WORDS} words.") | |
| chunks = split_chunks(text) | |
| if len(chunks) == 1: | |
| enc = tokenizer(_prompt(chunks[0]), return_tensors="pt", add_special_tokens=False).to(model.device) | |
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) | |
| Thread(target=model.generate, kwargs=dict(**enc, streamer=streamer, **_sampling(chunks))).start() | |
| out = "" | |
| for piece in streamer: | |
| out += piece | |
| yield out.split(STOP)[0].strip() | |
| yield out.split(STOP)[0].strip() | |
| return | |
| yield f"Rewriting {len(chunks)} chunks of about {CHUNK_WORDS} words in parallel..." | |
| enc = tokenizer([_prompt(c) for c in chunks], return_tensors="pt", padding=True, add_special_tokens=False).to(model.device) | |
| with torch.no_grad(): | |
| gen = model.generate(**enc, **_sampling(chunks)) | |
| texts = tokenizer.batch_decode(gen[:, enc["input_ids"].shape[1]:], skip_special_tokens=True) | |
| yield "\n\n".join(t.split(STOP)[0].strip() for t in texts) | |
| EXAMPLES = json.load(open(os.path.join(os.path.dirname(__file__), "examples.json"))) | |
| with gr.Blocks(title="Smaug Humanizer") as demo: | |
| gr.Markdown( | |
| "# Smaug Humanizer\n" | |
| "Paste AI-written text and get it back in a human voice, with the facts, numbers and layout kept. " | |
| "Model: [abacusai/Smaug-Humanizer](https://huggingface.co/abacusai/Smaug-Humanizer), a 12B open-weight model from Abacus.AI." | |
| ) | |
| with gr.Row(): | |
| inp = gr.Textbox(label="AI-written text", lines=16, max_lines=30, placeholder="Paste text here") | |
| out = gr.Textbox(label="Rewrite", lines=16, max_lines=30) | |
| btn = gr.Button("Humanize", variant="primary") | |
| gr.Examples(examples=[[e] for e in EXAMPLES], inputs=inp, label="Examples (chatbot answers on general topics)") | |
| btn.click(humanize, inputs=inp, outputs=out) | |
| demo.launch(mcp_server=True) | |