s123hree/green-code-optimizer-a100
0
1"""2inference.py — Inference endpoint for the Green-Code Optimizer 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 GRPO_OUTPUT_DIR – Training output root (default: ./grpo_output); Dockerfile sets /tmp/grpo_output on HF13"""14 15import os16import json17import torch18from peft import PeftModel19from transformers import AutoModelForCausalLM, AutoTokenizer20 21# ── Configuration ────────────────────────────────────────────────────────────22BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B-Instruct"23MAX_SEQ_LENGTH = 409624 25# Adapter source: HuggingFace Hub repo OR local path (must match train_grpo.py output layout)26HF_ADAPTER_REPO = os.getenv("HF_ADAPTER_REPO", "")27_repo_root = os.path.dirname(__file__)28_default_grpo_out = os.path.abspath(os.path.join(_repo_root, "grpo_output"))29_grpo_out = os.path.abspath(os.path.expanduser(os.environ.get("GRPO_OUTPUT_DIR", _default_grpo_out)))30ADAPTER_LOCAL_PATH = os.getenv(31 "ADAPTER_LOCAL_PATH",32 os.path.join(_grpo_out, "final_adapter"),33)34HF_TOKEN = os.getenv("HF_TOKEN", None)35 36SYSTEM_PROMPT = (37 "You are an expert Python refactoring agent focused on ENERGY EFFICIENCY.\n"38 "Your goal: minimise CPU cycles and peak memory while preserving program logic.\n"39 "Specifically prefer:\n"40 " • List/dict/set comprehensions over append-loops\n"41 " • Vectorised / built-in operations (sum, map) over manual accumulation\n"42 " • Hoisting loop-invariant work outside the loop\n"43 " • Eliminating dead code and redundant computation\n"44 " • Flattening unnecessarily nested loops\n"45 "Do NOT alter test files or break any existing assertions.\n"46 "Return edited files using EXACTLY this XML format:\n"47 '<file name="filename.py">\n... complete new code ...\n</file>\n'48 "Provide the full updated file content (do not omit any code)."49)50 51 52def _load_model():53 """Load the base model + LoRA adapter. Works with or without Unsloth."""54 device = "cuda" if torch.cuda.is_available() else "cpu"55 use_4bit = torch.cuda.is_available()56 57 # ── Try Unsloth first (faster, used during training) ─────────────────58 try:59 from unsloth import FastLanguageModel60 61 adapter_path = ADAPTER_LOCAL_PATH62 if HF_ADAPTER_REPO:63 adapter_path = HF_ADAPTER_REPO64 65 print(f"Loading via Unsloth: {adapter_path}")66 model, tokenizer = FastLanguageModel.from_pretrained(67 model_name=adapter_path,68 max_seq_length=MAX_SEQ_LENGTH,69 load_in_4bit=use_4bit,70 )71 FastLanguageModel.for_inference(model)72 if tokenizer.pad_token is None:73 tokenizer.pad_token = tokenizer.eos_token74 print("✅ Model loaded via Unsloth")75 return model, tokenizer76 77 except Exception as e:78 # Catch ImportError, runtime errors, OOM, etc. and fall back.79 print(f"Unsloth unavailable ({type(e).__name__}: {e}), "80 f"falling back to transformers+peft...")81 82 # ── Fallback: vanilla transformers + peft (Docker/CPU) ───────────────83 print(f"Loading base model: {BASE_MODEL}")84 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN)85 if tokenizer.pad_token is None:86 tokenizer.pad_token = tokenizer.eos_token87 88 model = AutoModelForCausalLM.from_pretrained(89 BASE_MODEL,90 torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,91 device_map="auto" if device == "cuda" else None,92 token=HF_TOKEN,93 )94 95 # Load adapter from HF Hub or local path96 adapter_source = HF_ADAPTER_REPO if HF_ADAPTER_REPO else ADAPTER_LOCAL_PATH97 if os.path.isdir(adapter_source) or HF_ADAPTER_REPO:98 print(f"Loading LoRA adapter: {adapter_source}")99 model = PeftModel.from_pretrained(model, adapter_source, token=HF_TOKEN)100 model = model.merge_and_unload()101 print("✅ LoRA adapter merged")102 else:103 print(f"⚠️ No adapter found at {adapter_source}, using base model")104 105 model.eval()106 print("✅ Model loaded via transformers+peft")107 return model, tokenizer108 109 110# ── Singleton model loading ──────────────────────────────────────────────────111_model = None112_tokenizer = None113 114 115def get_model():116 """Lazy-load the model on first call."""117 global _model, _tokenizer118 if _model is None:119 _model, _tokenizer = _load_model()120 return _model, _tokenizer121 122 123def run_inference(observation: dict) -> dict:124 """125 Given an environment observation, generate the next agent action.126 127 For the OpenEnv step-by-step interface (server.py):128 observation = {files, violation_report, steps_remaining, ...}129 130 Returns a tool-call dict: {"tool": "...", "args": {...}}131 """132 model, tokenizer = get_model()133 134 steps_remaining = observation.get("steps_remaining", 0)135 report = observation.get("violation_report", {})136 files = observation.get("files", {})137 138 # Build context from files139 code_context = ""140 for fname, content in files.items():141 block = f"\n--- {fname} ---\n```python\n{content}\n```\n"142 if len(code_context) + len(block) > 8000:143 break144 code_context += block145 146 user_prompt = (147 f"Steps remaining: {steps_remaining}\n"148 f"Violation report: {json.dumps(report, indent=2)}\n"149 f"\nHere is the codebase to refactor:\n{code_context}\n"150 f"\nPlease refactor and return the updated files."151 )152 153 messages = [154 {"role": "system", "content": SYSTEM_PROMPT},155 {"role": "user", "content": user_prompt},156 ]157 158 text_prompt = tokenizer.apply_chat_template(159 messages, tokenize=False, add_generation_prompt=True160 )161 inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)162 163 with torch.no_grad():164 outputs = model.generate(165 **inputs,166 max_new_tokens=512,167 do_sample=True,168 temperature=0.3,169 pad_token_id=tokenizer.pad_token_id,170 )171 172 generated = tokenizer.decode(173 outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True174 )175 176 # Parse the generated XML into edit actions177 import re178 file_pattern = r'<file name="(.*?)">(.*?)</file>'179 matches = re.findall(file_pattern, generated, flags=re.DOTALL)180 181 if matches:182 # Return the first edit as an action183 fname, content = matches[0]184 return {185 "tool": "edit_file",186 "args": {"filename": fname.strip(), "content": content.strip()},187 }188 189 # Fallback: check compliance190 return {"tool": "check_compliance", "args": {}}191 192 193# ── CLI entrypoint for quick testing ─────────────────────────────────────────194if __name__ == "__main__":195 print("=" * 60)196 print("🔍 Loading Model...")197 print("=" * 60)198 199 model, tokenizer = get_model()200 201 messy_code = '''202def calculate(x, y):203 import os204 a = x + y205 if a > 10:206 return a207 else:208 return 0209'''210 211 user_prompt = (212 f"Here is the codebase to refactor:\n\n"213 f"--- messy.py ---\n```python\n{messy_code}\n```\n\n"214 f"Please refactor and return the updated files."215 )216 messages = [217 {"role": "system", "content": SYSTEM_PROMPT},218 {"role": "user", "content": user_prompt},219 ]220 text_prompt = tokenizer.apply_chat_template(221 messages, tokenize=False, add_generation_prompt=True222 )223 inputs = tokenizer(text_prompt, return_tensors="pt").to(model.device)224 225 print("\n" + "=" * 60)226 print("🤖 Generating Refactored Code...")227 print("=" * 60)228 229 with torch.no_grad():230 outputs = model.generate(231 **inputs,232 max_new_tokens=512,233 do_sample=True,234 temperature=0.3,235 pad_token_id=tokenizer.pad_token_id,236 )237 238 generated_text = tokenizer.decode(239 outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True240 )241 print("--- AI OUTPUT ---")242 print(generated_text)243 print("-----------------")244