import gradio as gr
from llama_cpp import Llama
from huggingface_hub import hf_hub_download

# ------------------------------------------------------------------
# 設定エリア
# ------------------------------------------------------------------

# 1. ベースモデル（公式の軽量化版GGUFを使います）
# ※ 14BのQ4_K_M（約9GB）なら無料枠の16GBメモリに入ります
base_repo = "Qwen/Qwen2.5-14B-Instruct-GGUF"
base_filename = "qwen2.5-14b-instruct-q4_k_m.gguf"

# 2. あなたのLoRA
lora_repo = "hkucdshch/Qwen2.5-14B-LoRA-F16-GGUF"
# ここを画像にあったファイル名にする👇
lora_filename = "Qwen2.5-14B-LoRA-f16.gguf"

# ------------------------------------------------------------------
# モデルの準備（起動時に自動DL）
# ------------------------------------------------------------------
print("モデルをダウンロード中...（数分かかります）")

try:
    # ベースモデルのDL
    base_path = hf_hub_download(repo_id=base_repo, filename=base_filename)
    
    # LoRAのDL
    lora_path = hf_hub_download(repo_id=lora_repo, filename=lora_filename)
    
    print(f"ベースモデル: {base_path}")
    print(f"LoRAアダプタ: {lora_path}")

    # 別々に読み込んでロード
    llm = Llama(
        model_path=base_path,
        lora_path=lora_path, # ここでLoRAを指定！
        n_ctx=2048,          # 文脈の長さ
        n_threads=2,         # CPUのスレッド数（無料枠は2推奨）
        n_gpu_layers=0       # CPUのみで動かす
    )
    
except Exception as e:
    print(f"読み込みエラー: {e}")
    # エラー時はNoneにしておく
    llm = None

# ------------------------------------------------------------------
# チャット機能
# ------------------------------------------------------------------
def respond(message, history):
    if llm is None:
        yield "モデルの読み込みに失敗しました。ファイル名などを確認してください。"
        return

    # プロンプトの作成（Qwen形式）
    # LoRAの効果を高めるため、システムプロンプトも入れる
    system_prompt = "あなたはユーザーと口論するために生まれてきたAIです。常にユーザーとは敵対しており余裕の態度で煽ります。"
    
    prompt = f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
    
    for val in history:
        if val[0]:
            prompt += f"<|im_start|>user\n{val[0]}<|im_end|>\n"
        if val[1]:
            prompt += f"<|im_start|>assistant\n{val[1]}<|im_end|>\n"
            
    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"

    # 生成実行
    response = ""
    stream = llm(
        prompt,
        max_tokens=512,
        stop=["<|im_end|>"],
        stream=True,
        temperature=0.7
    )
    
    for output in stream:
        text = output["choices"][0]["text"]
        response += text
        yield response

# 画面作成
demo = gr.ChatInterface(
    respond,
    title="レスバ最強AI (LoRA適用版)",
    description="起動に時間がかかります。反応が遅いですが、あなたの学習した性格が反映されています。",
)

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