import astimport torchfrom transformers import AutoTokenizer, AutoModelForCausalLMfrom transformers.cache_utils import DynamicCache MODEL_ID = "kambleaa007/Self-Healing-LLM" tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)model = AutoModelForCausalLM.from_pretrained( MODEL_ID, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, device_map="auto" if torch.cuda.is_available() else "cpu") def trim_kv_cache(past_key_values: DynamicCache, target_len: int) -> None: """Crops the DynamicCache safely across all layers.""" if hasattr(past_key_values, "crop"): past_key_values.crop(target_len) else: for layer_idx in range(len(past_key_values.key_cache)): past_key_values.key_cache[layer_idx] = past_key_values.key_cache[layer_idx][:, :, :target_len, :] past_key_values.value_cache[layer_idx] = past_key_values.value_cache[layer_idx][:, :, :target_len, :] def is_valid_or_partial_code(code_str: str) -> bool: """Checks whether the code is valid or a valid partial prefix.""" try: ast.parse(code_str) return True except SyntaxError as e: if "unexpected EOF while parsing" in str(e) or "expected an indented block" in str(e): return True return False def generate_self_healing(prompt: str, max_tokens: int = 80, max_retries: int = 3): inputs = tokenizer(prompt, return_tensors="pt").to(model.device) token_stream = list(inputs.input_ids[0].cpu().numpy()) prompt_len = len(token_stream) past_key_values = DynamicCache() with torch.no_grad(): outputs = model(input_ids=inputs.input_ids, past_key_values=past_key_values, use_cache=True) past_key_values = outputs.past_key_values current_token = torch.argmax(outputs.logits[:, -1, :], dim=-1).unsqueeze(-1) line_checkpoints = [len(token_stream)] retry_count = 0 for _ in range(max_tokens): tid = current_token.item() token_text = tokenizer.decode([tid]) if "\n" in token_text: line_checkpoints.append(len(token_stream)) token_stream.append(tid) if "\n" in token_text and len(line_checkpoints) > 1: current_code = tokenizer.decode(token_stream, skip_special_tokens=True) if not is_valid_or_partial_code(current_code) and (retry_count < max_retries): retry_count += 1 last_valid_len = max(prompt_len, line_checkpoints[-2]) # Truncate KV Cache and rewind token stream trim_kv_cache(past_key_values, last_valid_len) token_stream = token_stream[:last_valid_len] line_checkpoints.pop() prev_token = torch.tensor([[token_stream[-1]]], device=model.device) with torch.no_grad(): redo_outputs = model(input_ids=prev_token, past_key_values=past_key_values, use_cache=True) past_key_values = redo_outputs.past_key_values # Sample probabilistically to break deterministic error loops scaled_logits = redo_outputs.logits[:, -1, :] / 0.8 current_token = torch.multinomial(torch.softmax(scaled_logits, dim=-1), num_samples=1) continue else: retry_count = 0 with torch.no_grad(): outputs = model(input_ids=current_token, past_key_values=past_key_values, use_cache=True) past_key_values = outputs.past_key_values current_token = torch.argmax(outputs.logits[:, -1, :], dim=-1).unsqueeze(-1) if tid == tokenizer.eos_token_id: break return tokenizer.decode(token_stream, skip_special_tokens=True) # Example Executioncode_prompt = "def calculate_discount(price, rate):\n # Return discounted price\n"print(generate_self_healing(code_prompt)) ## Training Details ### Training Procedure * Extended base tokenizer vocabulary with custom control tokens (`<UNDO>`, `<SEEK_LINE>`, `<EXEC_CHECKPOINT>`).* SFT fine-tuning on synthetic mutation-repair and diff-stream trajectories.* Evaluated on syntax-preserving line rollbacks using Hugging Face `transformers.Trainer` and `DataCollatorForLanguageModeling`. #### Training Hyperparameters * **Training regime:** fp16 mixed precision* **Learning rate:** 2e-4* **Optimizer:** AdamW* **Warmup steps:** 10 --- ## Technical Specifications ### Model Architecture and Objective * **Base Architecture:** Qwen2.5 (Decoder-only Transformer)* **Parameters:** ~494 Million* **Context Length:** 32,768 tokens* **Positional Embeddings:** RoPE (Rotary Position Embeddings)* **Attention Cache:** Paged dynamic KV-cache truncation compatible with `transformers.cache_utils.DynamicCache` --- ## Citation ```bibtex@misc{kamble2026selfhealingllm, author = {Ashish Kamble}, title = {Self-Healing-LLM: Non-Monotonic Decoding and In-Place KV-Cache Splicing for Code Synthesis}, year = {2026}, publisher = {Hugging Face}, howpublished = {\url{[https://huggingface.co/kambleaa007/Self-Healing-LLM](https://huggingface.co/kambleaa007/Self-Healing-LLM)}}}