Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model_id = "google/gemma-4-31B"
5quant_dir = "FilledVaccum/gemma-4-31B-awq-manual"
6
7# Load base model
8model = AutoModelForCausalLM.from_pretrained(
9 model_id,
10 torch_dtype=torch.float16,
11 device_map="auto",
12 trust_remote_code=True
13)
14tokenizer = AutoTokenizer.from_pretrained(quant_dir)
15
16# Load quantized weights
17checkpoint = torch.load(f"{quant_dir}/quantized_weights.pt", map_location="cpu")
18quantized_layers = checkpoint['quantized_layers']
19
20# Dequantize function
21def dequantize_layer(qdata, group_size=128):
22 weight = torch.zeros(qdata['shape'], dtype=torch.float16)
23 for i, scale_q in enumerate(qdata['qscales']):
24 start, end = i * group_size, min((i + 1) * group_size, qdata['shape'][1])
25 weight[:, start:end] = qdata['qweight'][:, start:end].float() * scale_q
26 return weight * qdata['scale_factor'].unsqueeze(0)
27
28# Load quantized weights into model
29for name, module in model.named_modules():
30 if name in quantized_layers:
31 module.weight.data = dequantize_layer(quantized_layers[name]).to(module.weight.device)
32
33# Run inference
34prompt = "What is artificial intelligence?"
35inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
36outputs = model.generate(**inputs, max_new_tokens=100)
37print(tokenizer.decode(outputs[0], skip_special_tokens=True))