import streamlit as st
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# 1. LOAD THE MODEL INTO THE SERVER'S RAM
@st.cache_resource
def load_model():
    # This securely downloads your GGUF file from your Hugging Face account
    # Make sure "llama-3-8b.Q4_K_M.gguf" matches your exact filename on Hugging Face
    model_path = hf_hub_download(
        repo_id="Aryanvaidh1712/AI_Humanizer-2", 
        filename="llama-3-8b.Q4_K_M.gguf" 
    )
    
    # Initialize the CPU inference engine
    llm = Llama(
        model_path=model_path,
        n_ctx=1024,      # Context window limit
        n_threads=8,     # Maximize the server's CPU cores
    )
    return llm

# 2. BUILD THE UI
st.set_page_config(page_title="AI Humanizer", page_icon="✨")
st.title("✨ AI Text Humanizer")

user_text = st.text_area("Original AI Text:", height=150)

# 3. GENERATION LOGIC
if st.button("Humanize Text"):
    if user_text:
        with st.spinner("The model is rewriting your text... (This takes a moment on free CPUs)"):
            llm = load_model()
            
            # The exact Alpaca prompt format from your training
            prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
Humanize the following text by converting it into active voice and adding natural transitions. Preserve meaning.
### Input:
{user_text}
### Response:
"""
            # Generate the text using your parameters
            output = llm(
                prompt,
                max_tokens=512,
                temperature=0.85,
                repeat_penalty=1.2,
                stop=["<|eot_id|>","<|end_of_text|>", "### Instruction:"],
                echo=False 
            )
            
            final_text = output["choices"][0]["text"].strip()
            
            st.success("Generation Complete!")
            st.write(final_text)
    else:
        st.warning("Please paste some text first.")