Update test.py

test.py

@@ -3,19 +3,15 @@ from transformers import AutoTokenizer
 import torch
 
 model_id = "LnL-AI/Yi-6B-gptq-4bit"
-
 device = torch.device("cuda:0")
 
-model_q = AutoGPTQForCausalLM.from_quantized(model_id, use_safetensors=True, device=device)
-
+model = AutoGPTQForCausalLM.from_quantized(model_id, use_safetensors=True, device=device)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 prompt = "I am in Shanghai, preparing to visit the natural history museum. Can you tell me the best way to"
+inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
-inp = tokenizer(prompt, return_tensors="pt").to(device)
-
-res = model_q.generate(**inp, num_beams=1, min_new_tokens=1, max_new_tokens=512)
+res = model.generate(**inputs, num_beams=1, min_new_tokens=1, max_new_tokens=512)
 
 predicted_text = tokenizer.decode(res[0])
-
 print(predicted_text)