# model_loader.py
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch

def load_model():
    # Define o modelo base e o caminho dos adapters (repositório atual)
    base_model = "defog/sqlcoder-7b-2"
    adapter_path = "./"  # Aqui, assume que os arquivos dos adapters estão no diretório raiz do repositório

    # Carregar o tokenizer
    tokenizer = AutoTokenizer.from_pretrained(adapter_path)
    tokenizer.pad_token = tokenizer.eos_token

    # Carregar o modelo base com quantização (assumindo 4-bit e utilização de fp16)
    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        device_map="auto",
        load_in_4bit=True,
        torch_dtype=torch.float16
    )
    model.config.pad_token_id = tokenizer.pad_token_id

    # Aplicar os adapters LoRA a partir do adapter_path
    model = PeftModel.from_pretrained(model, adapter_path)
    
    return model, tokenizer

if __name__ == "__main__":
    model, tokenizer = load_model()
    prompt = "portfolio_transaction_headers(...) JOIN portfolio_transaction_details(...): Find transactions for portfolio 72 involving LTC"
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=128)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))