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

BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B"
ADAPTER = "Nayan7r839/EchooaiForge"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

print("Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

print("Loading EchooaiForge adapter...")
model = PeftModel.from_pretrained(
    base_model,
    ADAPTER
)

model.eval()

print("EchooaiForge loaded!")


def generate_code(prompt: str):
    text = f"""### Instruction:
{prompt}

### Response:
"""

    inputs = tokenizer(
        text,
        return_tensors="pt"
    ).to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.2,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.05
        )

    result = tokenizer.decode(
        outputs[0],
        skip_special_tokens=True
    )

    if "### Response:" in result:
        result = result.split("### Response:", 1)[1]

    return result.strip()