import modal

APP_NAME = "pythia-410m-dolly-test"
VOLUME_NAME = "pythia-410m-dolly-output"

BASE_MODEL = "EleutherAI/pythia-410m"

VOLUME_ROOT = "/data"
ADAPTER_DIR = (
    f"{VOLUME_ROOT}/outputs/"
    "pythia-410m-dolly-lora/final-adapter"
)
CACHE_DIR = f"{VOLUME_ROOT}/cache/huggingface"

app = modal.App(APP_NAME)

volume = modal.Volume.from_name(
    VOLUME_NAME,
    create_if_missing=False,
)

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "torch==2.7.1",
        "transformers==4.53.2",
        "peft==0.16.0",
        "accelerate==1.8.1",
        "safetensors==0.5.3",
        "sentencepiece==0.2.0",
    )
    .env(
        {
            "HF_HOME": CACHE_DIR,
            "TRANSFORMERS_CACHE": CACHE_DIR,
            "TOKENIZERS_PARALLELISM": "false",
        }
    )
)


@app.function(
    image=image,
    gpu="T4",
    cpu=2.0,
    memory=8192,
    timeout=15 * 60,
    volumes={
        VOLUME_ROOT: volume.with_mount_options(read_only=True),
    },
)
def generate(
    instruction: str,
    context: str = "",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    top_p: float = 0.9,
    top_k: int = 40,
    repetition_penalty: float = 1.1,
    seed: int = 3407,
) -> str:
    import os

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

    if not os.path.isdir(ADAPTER_DIR):
        raise FileNotFoundError(
            f"Adapter directory was not found: {ADAPTER_DIR}"
        )

    required_files = [
        "adapter_config.json",
        "adapter_model.safetensors",
    ]

    for filename in required_files:
        path = os.path.join(ADAPTER_DIR, filename)
        if not os.path.isfile(path):
            raise FileNotFoundError(
                f"Required adapter file was not found: {path}"
            )

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA GPU was not detected.")

    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"Base model: {BASE_MODEL}")
    print(f"Adapter: {ADAPTER_DIR}")

    # The tokenizer was saved with the final LoRA adapter.
    tokenizer = AutoTokenizer.from_pretrained(
        ADAPTER_DIR,
        cache_dir=CACHE_DIR,
        use_fast=True,
    )

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

    tokenizer.padding_side = "left"

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

    model = PeftModel.from_pretrained(
        base_model,
        ADAPTER_DIR,
        is_trainable=False,
    )

    model = model.to("cuda")
    model.eval()

    clean_instruction = instruction.strip()
    clean_context = context.strip()

    if clean_context:
        prompt = (
            "### Instruction:\n"
            f"{clean_instruction}\n\n"
            "### Context:\n"
            f"{clean_context}\n\n"
            "### Response:\n"
        )
    else:
        prompt = (
            "### Instruction:\n"
            f"{clean_instruction}\n\n"
            "### Response:\n"
        )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        truncation=True,
        max_length=768,
    ).to("cuda")

    input_length = inputs["input_ids"].shape[1]

    generation_kwargs = {
        "max_new_tokens": max_new_tokens,
        "pad_token_id": tokenizer.pad_token_id,
        "eos_token_id": tokenizer.eos_token_id,
        "repetition_penalty": repetition_penalty,
        "use_cache": True,
    }

    # Temperature 0 means greedy decoding.
    if temperature > 0:
        generation_kwargs.update(
            {
                "do_sample": True,
                "temperature": temperature,
                "top_p": top_p,
                "top_k": top_k,
            }
        )
    else:
        generation_kwargs["do_sample"] = False

    with torch.inference_mode():
        output_ids = model.generate(
            **inputs,
            **generation_kwargs,
        )

    # Decode only tokens generated after the prompt.
    generated_ids = output_ids[0, input_length:]

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

    return response


@app.local_entrypoint()
def main(
    prompt: str = "Explain artificial intelligence in simple terms.",
    context: str = "",
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    top_p: float = 0.9,
    top_k: int = 40,
    repetition_penalty: float = 1.1,
    seed: int = 3407,
):
    response = generate.remote(
        instruction=prompt,
        context=context,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        repetition_penalty=repetition_penalty,
        seed=seed,
    )

    print("\n" + "=" * 72)
    print("PROMPT")
    print("=" * 72)
    print(prompt)

    if context:
        print("\n" + "=" * 72)
        print("CONTEXT")
        print("=" * 72)
        print(context)

    print("\n" + "=" * 72)
    print("MODEL RESPONSE")
    print("=" * 72)
    print(response)
    print("=" * 72)