Aniruddha7/QueryLens-Text2SQL_DocVQA-V2

1 2# Set environment variables BEFORE importing torch3import os4os.environ["TORCHINDUCTOR_DISABLE"] = "1"5os.environ["TORCH_COMPILE_DISABLE"] = "1"6# Optional: force a specific GPU (uncomment if needed)7# os.environ["CUDA_VISIBLE_DEVICES"] = "0"8 9from transformers import AutoModelForCausalLM, AutoTokenizer10from peft import PeftModel11import torch12 13# In-code disables for torch.compile and inductor/triton14torch._dynamo.config.suppress_errors = True15#torch._dynamo.config.enable = False16torch.set_float32_matmul_precision('high')17 18#base = "google/gemma-3-1b-it"19#adapter = "Agent/Gemma3_1B_weights"20base = "Qwen/Qwen2.5-3B-Instruct"21adapter = "Agent/Qwen2.5_3b_weights"22model = AutoModelForCausalLM.from_pretrained(23    base,24    dtype="auto",25    device_map=None,26    load_in_8bit=True,27    low_cpu_mem_usage=True,28    local_files_only=False  # Set to False to force download if needed29    #trust_remote_code=True30)31model = PeftModel.from_pretrained(model, adapter, local_files_only=True)32model = model.to("cuda")  # Move model to GPU33tokenizer = AutoTokenizer.from_pretrained(base)34 35prompt = "Schema: employee(emp_no, birth_date, first_name, last_name, gender, hire_date),\n" \36         "department(dept_no, dept_name),\n" \37         "dept_emp(emp_no, dept_no, from_date, to_date),\n" \38         "salary(emp_no, salary, from_date, to_date),\n" \39         "Write PostgreSQL query for highest and lowest salary across all departments. Do not include any explanation."40inputs = tokenizer(prompt, return_tensors="pt")41# Move all input tensors to CUDA42inputs = {k: v.to("cuda") for k, v in inputs.items()}43 44output = model.generate(**inputs, max_new_tokens=1028)45print(tokenizer.decode(output[0], skip_special_tokens=True))