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

# ------------------------------------------------------------------
# 設定エリア
# ------------------------------------------------------------------
# 1. ベースモデル
base_repo = "bartowski/Qwen2.5-14B-Instruct-GGUF"
base_filename = "Qwen2.5-14B-Instruct-Q4_K_M.gguf"

# 2. LoRA
lora_repo = "hkucdshch/portfolio6-FT"
lora_filename = "Qwen2.5-14B-TestLoRA-f16.gguf"

# ------------------------------------------------------------------
# モデルのロード
# ------------------------------------------------------------------
print(f"ベースモデル({base_filename})をダウンロード中...")

try:
    # ベースモデルのDL
    base_path = hf_hub_download(repo_id=base_repo, filename=base_filename)
    
    # LoRAのDL
    print(f"LoRA({lora_filename})をダウンロード中...")
    lora_path = hf_hub_download(repo_id=lora_repo, filename=lora_filename)
    
    print("モデルをメモリに読み込んでいます...")

    # ロード実行
    llm = Llama(
        model_path=base_path,
        lora_path=lora_path,
        lora_scale=1.0,
        n_ctx=1024,      # メモリ節約設定
        n_threads=2,     # CPUコア数
        n_gpu_layers=0,  # CPUのみ
        n_batch=512,     # メモリ節約
        verbose=True
    )
    print("【成功】準備完了！")
    
except Exception as e:
    print(f"【致命的エラー】: {e}")
    llm = None

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

    # システムプロンプト（性格設定）
    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"

    try:
        stream = llm(
            prompt,
            max_tokens=256,
            stop=["<|im_end|>"],
            stream=True,
            temperature=0.8,       # ランダムさ
            top_p=0.95,            # 幅広い表現を使わせる
            repeat_penalty=1.2,    # 同じような短文煽りを繰り返すのを防ぐ
            top_k=40,
        )
        
        response = ""
        for output in stream:
            text = output["choices"][0]["text"]
            response += text
            yield response
            
    except Exception as e:
        yield f"生成エラー: {str(e)}"

# 画面
demo = gr.ChatInterface(
    respond,
    title="ポートフォリオ用テスト口論AI (Qwen2.5-14B)",
    description="重いモデルのため返答までに数分待つことがあります。",
)

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