import gradio as gr
import torch
import ai_thingy


# Hugging Face Spaces can use CPU unless you've
# explicitly enabled a GPU.
ai_thingy.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


print("[SYSTEM] Initializing Aoban 1.1A-Refined...")
print(f"[SYSTEM] Device: {ai_thingy.device}")

ai_thingy.initialize_or_retrain(
    initial_train=True,
    use_live_data=False,
    epochs=0
)

print("[SYSTEM] Aoban 1.1A-Refined is ready.")


def respond(message, history, system_message, max_tokens, temperature, top_k):

    full_prompt = (
        f"{system_message}\n"
        f"User: {message}\n"
        f"Assistant:"
    )

    response = ai_thingy.generate_text(
        model=ai_thingy.current_model,
        tokenizer=ai_thingy.current_tokenizer,
        prompt=full_prompt,
        max_len=int(max_tokens),
        device=ai_thingy.device,
        top_k=int(top_k),
        penalty=1.8,
        temperature=float(temperature)
    )

    yield response


demo = gr.ChatInterface(
    fn=respond,
    title="AobanCorp™ Aoban Chatbot",
    description="Welcome to the official Aoban 1.1A-Refined chat.",
    additional_inputs=[
        gr.Textbox(
            value="You are a helpful but not really helpful AobanCorp™ assistant.",
            label="Aoban System Prompt"
        ),
        gr.Slider(
            minimum=1,
            maximum=31,
            value=20,
            step=1,
            label="Response Length"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=1.0,
            step=0.1,
            label="Creativity (Temperature)"
        ),
        gr.Slider(
            minimum=1,
            maximum=100,
            value=40,
            step=1,
            label="Focus (Top-K)"
        ),
    ],
)


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