| 76 | @pytest.mark.parametrize("dtype", [torch.float16], ids=describe_dtype) |
| 77 | @pytest.mark.slow |
| 78 | def test_pi(requires_cuda, model_and_tokenizer, inference_kernel, DQ, dtype): |
| 79 | fixture_config, model, tokenizer = model_and_tokenizer |
| 80 | |
| 81 | generation_config = transformers.GenerationConfig( |
| 82 | max_new_tokens=20, |
| 83 | do_sample=True, |
| 84 | top_p=0.9, |
| 85 | temperature=0.7, |
| 86 | ) |
| 87 | generation_config.max_new_tokens = 20 |
| 88 | |
| 89 | # text = 'Please write down the first 50 digits of pi.' |
| 90 | # text = get_prompt_for_generation_eval(text) |
| 91 | # text += ' Sure, here the first 50 digits of pi: 3.14159' |
| 92 | n_cases = 6 |
| 93 | text = "3.14159" |
| 94 | if hasattr(model.config, "quantization_config"): |
| 95 | model.config.quantization_config.bnb_4bit_compute_dtype = dtype |
| 96 | model.config.quantization_config.bnb_4bit_use_double_quant = DQ |
| 97 | |
| 98 | if not inference_kernel: |
| 99 | text = [text] * n_cases |
| 100 | inputs = tokenizer(text, return_tensors="pt").to("cuda:0") |
| 101 | x = inputs["input_ids"] |
| 102 | outputs = [] |
| 103 | if inference_kernel: |
| 104 | for i in range(n_cases): |
| 105 | output = model.generate(x, generation_config=generation_config) |
| 106 | textout = tokenizer.decode(output[0], skip_special_tokens=True) |
| 107 | outputs.append(textout) |
| 108 | else: |
| 109 | outputs = model.generate(x, generation_config=generation_config) |
| 110 | outputs = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs] |
| 111 | |
| 112 | assert len(outputs) == n_cases |
| 113 | failure_count = 0 |
| 114 | for i in range(n_cases): |
| 115 | if outputs[i][: len(str(math.pi))] != str(math.pi): |
| 116 | failure_count += 1 |
| 117 | failure_max = 2 if fixture_config[0] == "huggyllama/llama-7b" else 4 |
| 118 | if failure_count > failure_max: |
| 119 | print(math.pi) |
| 120 | for out in outputs: |
| 121 | print(out) |
| 122 | raise ValueError(f"Failure count: {failure_count}/{n_cases}") |