Initialize the JSON generator with a specific model.
(self, model_name: str = "Qwen/Qwen2.5-Coder-0.5B-Instruct")
| 23 | return torch.device("cpu") |
| 24 | |
| 25 | def __init__(self, model_name: str = "Qwen/Qwen2.5-Coder-0.5B-Instruct"): |
| 26 | """Initialize the JSON generator with a specific model.""" |
| 27 | self.device = self.get_device() |
| 28 | logger.info(f"Using device: {self.device}") |
| 29 | try: |
| 30 | # Initialize the model and tokenizer using the new outlines API |
| 31 | hf_model = AutoModelForCausalLM.from_pretrained( |
| 32 | model_name, |
| 33 | device_map="auto" if str(self.device) != "cpu" else None, |
| 34 | torch_dtype=torch.float16 if str(self.device) != "cpu" else torch.float32 |
| 35 | ) |
| 36 | self.tokenizer = AutoTokenizer.from_pretrained(model_name) |
| 37 | |
| 38 | # Create outlines model |
| 39 | self.model = outlines.from_transformers(hf_model, self.tokenizer) |
| 40 | logger.info(f"Successfully loaded model: {model_name}") |
| 41 | except Exception as e: |
| 42 | logger.error(f"Error loading model: {str(e)}") |
| 43 | raise |
| 44 | |
| 45 | def count_tokens(self, text: str) -> int: |
| 46 | """Count the number of tokens in a text string.""" |
nothing calls this directly
no test coverage detected