import gc
import torch
import gradio as gr
from transformers import AutoProcessor, AutoModelForImageTextToText


# ============================================================
# 模型設定
# ============================================================
MODEL_ID = "google/gemma-4-12B"


# ============================================================
# 清理記憶體
# ============================================================
gc.collect()

if torch.cuda.is_available():
    torch.cuda.empty_cache()


# ============================================================
# 載入模型與 Processor
# ============================================================
print(f"Loading model: {MODEL_ID}")

processor = AutoProcessor.from_pretrained(MODEL_ID)

model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    low_cpu_mem_usage=True,
)

model.eval()

print("Model loaded successfully.")


# ============================================================
# 產生回答
# ============================================================
def generate_gemma(message, history):
    tokenizer = processor.tokenizer

    # 建立簡單對話 prompt
    prompt_parts = []

    if history:
        for user_msg, assistant_msg in history:
            if user_msg:
                prompt_parts.append(f"使用者：{user_msg}")
            if assistant_msg:
                prompt_parts.append(f"助理：{assistant_msg}")

    prompt_parts.append(f"使用者：{message}")
    prompt_parts.append("助理：")

    prompt = "\n".join(prompt_parts)

    inputs = processor(
        text=prompt,
        return_tensors="pt",
    )

    inputs = {
        key: value.to(model.device)
        for key, value in inputs.items()
    }

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id,
        )

    input_length = inputs["input_ids"].shape[-1]
    generated_ids = outputs[0][input_length:]

    response = tokenizer.decode(
        generated_ids,
        skip_special_tokens=True,
    ).strip()

    return response


# ============================================================
# Gradio 介面
# ============================================================
demo = gr.ChatInterface(
    fn=generate_gemma,
    title="Gemma 4-12B Chat",
    description="使用 Gemma 4-12B 的 Gradio 推論介面",
    examples=[
        "請使用繁體中文說明什麼是內部稽核？",
        "內部稽核與外部稽核有何不同？",
        "請說明內部控制的五大要素。",
        "什麼是三道防線模型？",
    ],
)


# ============================================================
# 啟動 App
# ============================================================
if __name__ == "__main__":
    demo.launch()