# Compatibility fix
import huggingface_hub

if not hasattr(huggingface_hub, "HfFolder"):
    class HfFolder:
        @staticmethod
        def get_token():
            return huggingface_hub.get_token()

        @staticmethod
        def save_token(token):
            return huggingface_hub.login(token=token)

        @staticmethod
        def delete_token():
            try:
                huggingface_hub.logout()
            except Exception:
                pass

    huggingface_hub.HfFolder = HfFolder


import spaces
import gradio as gr
import torch

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig


ADAPTER = "lsgz/lsgz-personality-clone"

# Get base model from your LoRA config
config = PeftConfig.from_pretrained(ADAPTER)
BASE_MODEL = config.base_model_name_or_path

print("Base model:", BASE_MODEL)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token


# -------------------------
# LOAD MODEL ON CPU
# -------------------------

print("Loading base model...")

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
)

print("Loading LSGZ adapter...")

model = PeftModel.from_pretrained(
    base_model,
    ADAPTER
)

model.eval()

print("Model ready on CPU.")


# -------------------------
# GPU INFERENCE
# -------------------------

@spaces.GPU(duration=120)
def respond(message, history):

    print("GPU available:", torch.cuda.is_available())
    print("GPU:", torch.cuda.get_device_name(0))

    # GPU exists HERE
    model.to("cuda")

    inputs = tokenizer(
        message,
        return_tensors="pt"
    ).to("cuda")

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

    generated = outputs[0][inputs["input_ids"].shape[1]:]

    response = tokenizer.decode(
        generated,
        skip_special_tokens=True
    )

    return response.strip()


# -------------------------
# GRADIO
# -------------------------

demo = gr.ChatInterface(
    fn=respond,
    title="LSGZ Personality Clone",
    description="Chat with LSGZ 💬",
)

demo.queue()
demo.launch()