Create model_loader.py

model_loader.py

@@ -0,0 +1,34 @@
+# 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))