app.py

import os
import re
import torch
import gradio as gr
from threading import Thread
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TextIteratorStreamer,
)
from peft import PeftModel

BASE_MODEL = "Qwen/Qwen2.5-14B-Instruct"
ADAPTER_NAME = "mutij8/academic-humanizer-qwen-14b"

print("⏳ Loading Tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_NAME)

print("⏳ Configuring 4-bit Quantization for Nvidia T4 (16GB VRAM)...")
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16,
)

print("⏳ Loading Base Model in 4-bit onto T4 GPU...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=bnb_config,
    device_map="auto",
    low_cpu_mem_usage=True,
)

print("⏳ Attaching Fine-Tuned LoRA Adapter...")
model = PeftModel.from_pretrained(base_model, ADAPTER_NAME)
model.eval()

print("⚡ Compiling model with torch.compile() for accelerated token generation...")
# Optimizes the execution graph for faster inference throughput on T4
try:
    model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
except Exception as e:
    print(f"⚠️ torch.compile() skipped or failed: {e}. Falling back to standard execution.")

print("✅ Model, Adapter & Compilation ready on Nvidia T4!")

SYSTEM_PROMPT = (
    "You are an expert human academic writer. Rewrite the provided AI text into natural, human-written prose "
    "that bypasses AI detection.\n"
    "STRICT RULES:\n"
    "1. Paraphrase ONLY the provided paragraphs. Never add extra introductions, placeholders ([Author], [Year]), or LaTeX code.\n"
    "2. Preserve every inline citation (e.g., (Smith, 2021)) and section heading accurately.\n"
    "3. Use natural human sentence variation, active/passive voice shifts, and varied clause lengths."
)

def humanize_full_essay(essay_text, temperature):
    if not essay_text or not essay_text.strip():
        yield ""
        return

    device = "cuda" if torch.cuda.is_available() else "cpu"
    paragraphs = [p.strip() for p in re.split(r'\n\s*\n', essay_text) if p.strip()]
    rewritten_paragraphs = []

    for p in paragraphs:
        if len(p.split()) <= 8 and not p.endswith("."):
            rewritten_paragraphs.append(p)
            yield "\n\n".join(rewritten_paragraphs)
            continue

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": p}
        ]

        inputs = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
        ).to(device)

        streamer = TextIteratorStreamer(
            tokenizer, skip_prompt=True, skip_special_tokens=True
        )
        word_count = len(p.split())

        gen_kwargs = dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=int(word_count * 2.2) + 80,
            temperature=float(temperature),
            top_p=0.9,
            repetition_penalty=1.12,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

        thread = Thread(target=model.generate, kwargs=gen_kwargs)
        thread.start()

        partial = ""
        prefix = "\n\n".join(rewritten_paragraphs)
        for token_text in streamer:
            partial += token_text
            yield f"{prefix}\n\n{partial}".strip() if prefix else partial.strip()

        rewritten_paragraphs.append(partial.strip())
        
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

demo = gr.Interface(
    fn=humanize_full_essay,
    inputs=[
        gr.Textbox(lines=16, label="Paste Full AI Essay"),
        gr.Slider(minimum=0.3, maximum=1.0, value=0.75, step=0.05, label="Temperature"),
    ],
    outputs=gr.Textbox(lines=16, label="Humanized Output (Nvidia T4 - Compiled)"),
    title="Academic Humanizer (Qwen2.5-14B Fine-Tuned)",
)

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)