| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| tokenizer = AutoTokenizer.from_pretrained("stabilityai/stable-code-3b") | |
| model = AutoModelForCausalLM.from_pretrained( | |
| "stabilityai/stable-code-3b", | |
| torch_dtype="auto", | |
| ) | |
| codeSnip = """ | |
| Is the following CSS code valid? | |
| ``` | |
| div { | |
| color: red; | |
| width: 100px; | |
| } | |
| ``` | |
| """ | |
| inputs = tokenizer(codeSnip, return_tensors="pt") | |
| tokens = model.generate( | |
| **inputs, | |
| max_new_tokens=48, | |
| temperature=0.2, | |
| do_sample=True, | |
| ) | |
| res = tokenizer.decode(tokens[0], skip_special_tokens=True) | |
| print(res) | |