razak123/code-migration-env

1# code_migration_env\inference_nvidia.py2 3import os4os.environ["OPENBLAS_NUM_THREADS"] = "1"5os.environ["OMP_NUM_THREADS"] = "1"6import re7import json8import asyncio9from typing import List, Optional, Dict, Tuple10 11from dotenv import load_dotenv12from openai import OpenAI13 14load_dotenv()15 16def get_env_var(name: str, required: bool = True, default: Optional[str] = None) -> str:17    value = os.environ.get(name, default)18    if required and not value:19        print(f"ERROR: Mandatory environment variable '{name}' is not set.", flush=True)20        raise SystemExit(1)21    return value or ""22 23NVIDIA_BASE_URL = get_env_var("NVIDIA_BASE_URL", required=False, default="https://integrate.api.nvidia.com/v1")24 25TASK_NAME = os.environ.get("TASK_NAME", "python_modernize")26IMAGE_NAME = os.environ.get("IMAGE_NAME", "huggingface/spaces/razak123/code-migration-env")27MAX_STEPS = int(os.environ.get("MAX_STEPS", "3"))28MAX_TOTAL_REWARD = float(os.environ.get("MAX_TOTAL_REWARD", "1.0"))29SUCCESS_SCORE_THRESHOLD = float(os.environ.get("SUCCESS_SCORE_THRESHOLD", "0.5"))30 31TASKS = [32    ("python_modernize", "easy"),33    ("python_to_node",   "medium"),34    ("pandas_to_polars", "hard"),35]36 37def log_start(task: str, env: str, model: str):38    print(f"[START] task={task} env={env} model={model}", flush=True)39 40def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str] = None):41    error_str = f" error={error}" if error else ""42    print(f"[STEP] step={step} action={action!r} reward={reward:.3f} done={done}{error_str}", flush=True)43 44def log_end(success: bool, steps: int, score: float, rewards: List[float]):45    print(f"[END] success={success} steps={steps} score={score:.3f} rewards={rewards}", flush=True)46 47def get_nvidia_model_candidates() -> List[Tuple[str, str]]:48    """49    Returns (model_name, api_key) pairs in preference order.50    Only non-empty entries are included.51    """52    candidates = [53        # (os.environ.get("NVIDIA_MODEL_STEP_FLASH", ""), os.environ.get("NVIDIA_KEY_STEP_FLASH", "")),54        (os.environ.get("NVIDIA_MODEL_DEVSTRAL", ""), os.environ.get("NVIDIA_KEY_DEVSTRAL", "")),55        # (os.environ.get("NVIDIA_MODEL_KIMI_K2", ""), os.environ.get("NVIDIA_KEY_KIMI_K2", "")),56        # (os.environ.get("NVIDIA_MODEL_MISTRAL_LARGE", ""), os.environ.get("NVIDIA_KEY_MISTRAL_LARGE", "")),57        # (os.environ.get("NVIDIA_MODEL_DEEPSEEK_V3_1", ""), os.environ.get("NVIDIA_KEY_DEEPSEEK_V3_1", "")),58        # (os.environ.get("NVIDIA_MODEL_MAGISTRAL_SMALL", ""), os.environ.get("NVIDIA_KEY_MAGISTRAL_SMALL", "")),59        # (os.environ.get("NVIDIA_MODEL_GLM47", ""), os.environ.get("NVIDIA_KEY_GLM47", "")),60        # (os.environ.get("gemma_4b_model", ""), os.environ.get("gemma_4b_key", ""))61    ]62    return [(model, key) for model, key in candidates if model and key]63 64def clean_json_response(text: str) -> str:65    text = text.strip()66    text = re.sub(r"^```json\s*", "", text, flags=re.IGNORECASE)67    text = re.sub(r"^```\s*", "", text)68    text = re.sub(r"```$", "", text)69    return text.strip()70 71def extract_json_object(text: str) -> Optional[Dict]:72    cleaned = clean_json_response(text)73    try:74        return json.loads(cleaned)75    except Exception as e:76        print(f"[DEBUG] JSON parse failed: {e}", flush=True)  # ADD77 78    # fallback: try to grab the first {...} block79    match = re.search(r"\{.*\}", cleaned, flags=re.DOTALL)80    if match:81        try:82            return json.loads(match.group(0))83        except Exception:84            return None85    return None86 87def build_prompt(obs) -> str:88    history = "\n".join(obs.history or [])89    return f"""90You are a code migration expert.91 92Migrate the following {obs.source_language} code to {obs.target_language}.93 94Task ID: {obs.task_id}95Difficulty: {obs.difficulty}96 97Requirements:98{obs.requirements}99 100Test description:101{obs.test_description}102 103Previous attempts:104{history if history else "(none)"}105 106Source code:107```{obs.source_language}108{obs.source_code}109 110Return ONLY valid JSON with exactly these keys:111 112translated_code113explanation114 115Do not include markdown fences.116""".strip()117 118async def call_model_for_action(model: str, api_key: str, prompt: str) -> Dict[str, str]:119    client = OpenAI(120        base_url=NVIDIA_BASE_URL,121        api_key=api_key,122    )123 124    try:125        completion = client.chat.completions.create(126            model=model,127            temperature=0.2,128            max_tokens=2048,129            messages=[130                {131                    "role": "system",132                    "content": "Return only valid JSON with keys translated_code and explanation."133                },134                {"role": "user", "content": prompt},135            ],136        )137 138        text = (completion.choices[0].message.content or "").strip()139        print(f"[DEBUG] Raw model response: {text[:200]}", flush=True)140        parsed = extract_json_object(text)141        print(f"[DEBUG] Parsed result: {parsed}", flush=True)142        if parsed and isinstance(parsed, dict):143            translated_code = str(parsed.get("translated_code", "")).strip()144            explanation = str(parsed.get("explanation", "")).strip()145            if translated_code:146                return {147                    "translated_code": translated_code,148                    "explanation": explanation or "NVIDIA-generated translation.",149                }150        151        # fallback: wrap raw output152        return {153            "translated_code": text,154            "explanation": "Raw model output could not be parsed as JSON.",155        }156    except Exception as e:157        print(f"[DEBUG] NVIDIA call failed for {model}: {e}", flush=True)158        return {159            "translated_code": "",160            "explanation": f"API error: {e}"161        }162 163async def run_single_task_with_env(env, task_name: str, episode_id: str, models, current_model):164    log_start(task=task_name, env=IMAGE_NAME, model=current_model)165 166    from models import CodeMigrationAction167 168    rewards: List[float] = []169    steps_taken = 0170    score = 0.0171    success = False172 173    try:174        # Reset with episode_id to switch scenario — no new container needed175        result = await env.reset(episode_id=episode_id)176 177        for step in range(1, MAX_STEPS + 1):178            if result.done:179                break180 181            obs = result.observation182            prompt = build_prompt(obs)183 184            action_data = None185            for model_name, api_key in models:186                res = await call_model_for_action(model_name, api_key, prompt)187                if res["translated_code"]:188                    action_data = res189                    current_model = model_name190                    break191 192            if not action_data:193                action_data = {"translated_code": "error", "explanation": "All models failed"}194 195            try:196                action = CodeMigrationAction(197                    translated_code=action_data["translated_code"],198                    explanation=action_data["explanation"][:1999]199                )200            except Exception as e:201                action = CodeMigrationAction(202                    translated_code="error",203                    explanation=f"Validation error: {e}"204                )205 206            result = await env.step(action)207            reward = result.reward or 0.0208            rewards.append(reward)209            steps_taken = step210 211            log_step(212                step=step,213                action=action.translated_code[:100] + ("..." if len(action.translated_code) > 100 else ""),214                reward=reward,215                done=result.done,216                error=None217            )218 219            if result.done:220                break221 222        score = sum(rewards) / MAX_TOTAL_REWARD if MAX_TOTAL_REWARD > 0 else 0.0223        score = min(max(score, 0.0), 1.0)224        success = score >= SUCCESS_SCORE_THRESHOLD225 226    except Exception as e:227        print(f"[ERROR] Task {task_name} failed: {e}", flush=True)228 229    log_end(success=success, steps=steps_taken, score=score, rewards=rewards)230    return {231        "task": task_name,232        "success": success,233        "steps": steps_taken,234        "score": score,235        "rewards": rewards236    }237 238 239async def main() -> None:240    models = get_nvidia_model_candidates()241    if not models:242        print("ERROR: No NVIDIA model/key pairs found.", flush=True)243        raise SystemExit(1)244 245    current_model = models[0][0]246    task_results = []247 248    from client import CodeMigrationEnv249    from models import CodeMigrationAction250 251    # ONE env for all tasks252    env = await CodeMigrationEnv.from_docker_image(IMAGE_NAME)253 254    try:255        for task_name, episode_id in TASKS:256            result = await run_single_task_with_env(257                env, task_name, episode_id, models, current_model258            )259            task_results.append(result)260    finally:261        try:262            await env.close()263        except Exception as e:264            print(f"[DEBUG] env.close() error: {e}", flush=True)265 266    print("\n[SUMMARY] Task Results:", flush=True)267    for r in task_results:268        print(f"  {r['task']}: success={r['success']} score={r['score']:.3f}", flush=True)269 270if __name__ == "__main__":271    asyncio.run(main())272 273