| 13 | logger = logging.getLogger(__name__) |
| 14 | |
| 15 | class JSONGenerator: |
| 16 | def get_device(self): |
| 17 | """Get the appropriate device (mps, cuda, or cpu).""" |
| 18 | if torch.backends.mps.is_available(): |
| 19 | return torch.device("mps") |
| 20 | elif torch.cuda.is_available(): |
| 21 | return torch.device("cuda") |
| 22 | else: |
| 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.""" |
| 47 | try: |
| 48 | tokens = self.tokenizer.encode(text) |
| 49 | return len(tokens) |
| 50 | except Exception as e: |
| 51 | logger.error(f"Error counting tokens: {str(e)}") |
| 52 | return 0 |
| 53 | |
| 54 | def parse_json_schema_to_pydantic(self, schema_str: str) -> type[BaseModel]: |
| 55 | """Convert JSON schema string to Pydantic model.""" |
| 56 | try: |
| 57 | schema_dict = json.loads(schema_str) |
| 58 | |
| 59 | # Extract properties and required fields |
| 60 | properties = schema_dict.get('properties', {}) |
| 61 | required = schema_dict.get('required', []) |
| 62 | |
| 63 | # Build field definitions for Pydantic |
| 64 | fields = {} |
| 65 | for field_name, field_def in properties.items(): |
| 66 | field_type = str # Default to string |
| 67 | |
| 68 | # Map JSON schema types to Python types |
| 69 | if field_def.get('type') == 'integer': |
| 70 | field_type = int |
| 71 | elif field_def.get('type') == 'number': |
| 72 | field_type = float |