import os
from huggingface_hub import login
login(token=os.getenv("HF_TOKEN"))

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    torch_dtype="auto"
)

generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_length=512,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.1
)

def humanize_text(text):
    if not text.strip():
        return "⚠️ Please enter some text."
    
    prompt = f"""Rewrite the following text to sound natural, fluent, and human-like. 
Preserve meaning, names, and numbers. Avoid robotic tone. 
Use contractions, natural sentence flow, and varied structure. 
Do not explain, only rewrite.

Input: \"\"\"{text}\"\"\"
Rewritten:"""
    
    output = generator(prompt, num_return_sequences=1)[0]["generated_text"]
    # Strip off prompt echo if model repeats
    if "Rewritten:" in output:
        output = output.split("Rewritten:")[-1].strip()
    return output

demo = gr.Interface(
    fn=humanize_text,
    inputs=gr.Textbox(lines=6, placeholder="Paste your text here..."),
    outputs=gr.Textbox(label="Humanized Output"),
    title="AI Humanizer",
    description="Drop text and get a more natural, human-like version. Powered by Mistral-7B-Instruct."
)

if __name__ == "__main__":
    demo.launch()