File size: 1,344 Bytes
5a17665 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | # 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))
|