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

model_id = "TheBloke/MythoMax-L2-13B-GPTQ"

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    trust_remote_code=True,
    revision="main",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

def chat(prompt):
    output = pipe(prompt, max_new_tokens=400, temperature=0.7, top_p=0.9, repetition_penalty=1.1)
    return output[0]["generated_text"]

gr.Interface(fn=chat,
             inputs=gr.Textbox(label="Prompt", lines=6, placeholder="Tulis kode atau pertanyaan..."),
             outputs=gr.Textbox(label="Respon MythoMax"),
             title="🧙‍♂️ MythoMax L2 13B Coder",
             description="Model LLM roleplay + coding kelas berat 🤘 oleh King Hammz"
            ).launch()