onlycoding135/constrained-refactor-gauntlet

1"""2inference.py — Inference endpoint for the Constrained Refactor Gauntlet agent.3 4Supports two modes:5  1. Local: loads adapter from grpo_output/final_adapter (with Unsloth)6  2. Docker/HF Spaces: downloads adapter from HuggingFace Hub (vanilla transformers+peft)7 8Environment variables:9  HF_ADAPTER_REPO  – HuggingFace repo ID for the adapter (e.g. "username/adapter-name")10  HF_TOKEN          – HuggingFace token for private repos11  ADAPTER_LOCAL_PATH – Override local path to adapter directory12"""13 14import os15import json16import torch17from peft import PeftModel18from transformers import AutoModelForCausalLM, AutoTokenizer19 20# ── Configuration ────────────────────────────────────────────────────────────21BASE_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"22MAX_SEQ_LENGTH = 409623 24# Adapter source: HuggingFace Hub repo OR local path25HF_ADAPTER_REPO = os.getenv("HF_ADAPTER_REPO", "")26ADAPTER_LOCAL_PATH = os.getenv(27    "ADAPTER_LOCAL_PATH",28    os.path.join(os.path.dirname(__file__), "grpo_output", "final_adapter")29)30HF_TOKEN = os.getenv("HF_TOKEN", None)31 32SYSTEM_PROMPT = (33    "You are an expert Python refactoring agent. Your task is to clean up the provided codebase, "34    "improve its quality (tests, linting, complexity), and fix compliance issues.\n"35    "You must return your edited files using the following exact XML format:\n"36    '<file name="filename.py">\n... complete new code ...\n</file>\n'37    "Do not omit any code inside the file block. Provide the full updated file."38)39 40 41def _load_model():42    """Load the base model + LoRA adapter. Works with or without Unsloth."""43    device = "cuda" if torch.cuda.is_available() else "cpu"44    use_4bit = torch.cuda.is_available()45 46    # ── Try Unsloth first (faster, used during training) ─────────────────47    try:48        from unsloth import FastLanguageModel49 50        adapter_path = ADAPTER_LOCAL_PATH51        if HF_ADAPTER_REPO:52            adapter_path = HF_ADAPTER_REPO53 54        print(f"Loading via Unsloth: {adapter_path}")55        model, tokenizer = FastLanguageModel.from_pretrained(56            model_name=adapter_path,57            max_seq_length=MAX_SEQ_LENGTH,58            load_in_4bit=use_4bit,59        )60        FastLanguageModel.for_inference(model)61        if tokenizer.pad_token is None:62            tokenizer.pad_token = tokenizer.eos_token63        print("✅ Model loaded via Unsloth")64        return model, tokenizer65 66    except (ImportError, Exception) as e:67        print(f"Unsloth unavailable ({e}), falling back to transformers+peft...")68 69    # ── Fallback: vanilla transformers + peft (Docker/CPU) ───────────────70    print(f"Loading base model: {BASE_MODEL}")71    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN)72    if tokenizer.pad_token is None:73        tokenizer.pad_token = tokenizer.eos_token74 75    model = AutoModelForCausalLM.from_pretrained(76        BASE_MODEL,77        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,78        device_map="auto" if device == "cuda" else None,79        token=HF_TOKEN,80    )81 82    # Load adapter from HF Hub or local path83    adapter_source = HF_ADAPTER_REPO if HF_ADAPTER_REPO else ADAPTER_LOCAL_PATH84    if os.path.isdir(adapter_source) or HF_ADAPTER_REPO:85        print(f"Loading LoRA adapter: {adapter_source}")86        model = PeftModel.from_pretrained(model, adapter_source, token=HF_TOKEN)87        model = model.merge_and_unload()88        print("✅ LoRA adapter merged")89    else:90        print(f"⚠️  No adapter found at {adapter_source}, using base model")91 92    model.eval()93    print("✅ Model loaded via transformers+peft")94    return model, tokenizer95 96 97# ── Singleton model loading ──────────────────────────────────────────────────98_model = None99_tokenizer = None100 101 102def get_model():103    """Lazy-load the model on first call."""104    global _model, _tokenizer105    if _model is None:106        _model, _tokenizer = _load_model()107    return _model, _tokenizer108 109 110def run_inference(observation: dict) -> dict:111    """112    Given an environment observation, generate the next agent action.113    114    For the OpenEnv step-by-step interface (server.py):115      observation = {files, violation_report, steps_remaining, ...}116    117    Returns a tool-call dict: {"tool": "...", "args": {...}}118    """119    model, tokenizer = get_model()120 121    steps_remaining = observation.get("steps_remaining", 0)122    report = observation.get("violation_report", {})123    files = observation.get("files", {})124 125    # Build context from files126    code_context = ""127    for fname, content in files.items():128        block = f"\n--- {fname} ---\n```python\n{content}\n```\n"129        if len(code_context) + len(block) > 8000:130            break131        code_context += block132 133    user_prompt = (134        f"Steps remaining: {steps_remaining}\n"135        f"Violation report: {json.dumps(report, indent=2)}\n"136        f"\nHere is the codebase to refactor:\n{code_context}\n"137        f"\nPlease refactor and return the updated files."138    )139 140    messages = [141        {"role": "system", "content": SYSTEM_PROMPT},142        {"role": "user", "content": user_prompt},143    ]144 145    text_prompt = tokenizer.apply_chat_template(146        messages, tokenize=False, add_generation_prompt=True147    )148    inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)149 150    with torch.no_grad():151        outputs = model.generate(152            **inputs,153            max_new_tokens=512,154            do_sample=True,155            temperature=0.3,156            pad_token_id=tokenizer.pad_token_id,157        )158 159    generated = tokenizer.decode(160        outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True161    )162 163    # Parse the generated XML into edit actions164    import re165    file_pattern = r'<file name="(.*?)">(.*?)</file>'166    matches = re.findall(file_pattern, generated, flags=re.DOTALL)167 168    if matches:169        # Return the first edit as an action170        fname, content = matches[0]171        return {172            "tool": "edit_file",173            "args": {"filename": fname.strip(), "content": content.strip()},174        }175 176    # Fallback: check compliance177    return {"tool": "check_compliance", "args": {}}178 179 180# ── CLI entrypoint for quick testing ─────────────────────────────────────────181if __name__ == "__main__":182    print("=" * 60)183    print("🔍 Loading Model...")184    print("=" * 60)185 186    model, tokenizer = get_model()187 188    messy_code = '''189def calculate(x, y):190    import os191    a = x + y192    if a > 10:193        return a194    else:195        return 0196'''197 198    user_prompt = (199        f"Here is the codebase to refactor:\n\n"200        f"--- messy.py ---\n```python\n{messy_code}\n```\n\n"201        f"Please refactor and return the updated files."202    )203    messages = [204        {"role": "system", "content": SYSTEM_PROMPT},205        {"role": "user", "content": user_prompt},206    ]207    text_prompt = tokenizer.apply_chat_template(208        messages, tokenize=False, add_generation_prompt=True209    )210    inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)211 212    print("\n" + "=" * 60)213    print("🤖 Generating Refactored Code...")214    print("=" * 60)215 216    with torch.no_grad():217        outputs = model.generate(218            **inputs,219            max_new_tokens=512,220            do_sample=True,221            temperature=0.3,222            pad_token_id=tokenizer.pad_token_id,223        )224 225    generated_text = tokenizer.decode(226        outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True227    )228    print("--- AI OUTPUT ---")229    print(generated_text)230    print("-----------------")231