Create models/loader.py

models/loader.py

@@ -0,0 +1,45 @@
+# models/loader.py
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+
+MODEL_REGISTRY = {
+    "ceo": "Qwen/Qwen3-0.6B",
+    "manager": "Qwen/Qwen3-0.6B",
+    "worker_coder": "Qwen/Qwen3-0.6B",
+    "worker_tester": "Qwen/Qwen3-0.6B",
+}
+_MODEL_CACHE = {}
+
+def load_model(model_name):
+    if model_name in _MODEL_CACHE:
+        return _MODEL_CACHE[model_name]
+    
+    model_kwargs = {"device_map": "auto", "trust_remote_code": True, "attn_implementation": "eager"}
+    if torch.cuda.is_available():
+        print(f"CUDA found. Loading '{model_name}' in 4-bit.")
+        bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16)
+        model_kwargs["quantization_config"] = bnb_config
+    else:
+        print(f"CUDA not found. Loading '{model_name}' on CPU.")
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
+    _MODEL_CACHE[model_name] = (tokenizer, model)
+    print(f"Model {model_name} loaded and cached.")
+    return tokenizer, model
+
+def generate_with_model(role: str, prompt: str) -> str:
+    from backend.agents import ROLE_PROMPTS
+    from .loader import MODEL_REGISTRY
+    
+    try:
+        model_name = MODEL_REGISTRY[role]
+        tokenizer, model = load_model(model_name)
+        messages = [{"role": "system", "content": ROLE_PROMPTS[role]}, {"role": "user", "content": prompt}]
+        input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=3072, pad_token_id=tokenizer.eos_token_id, use_cache=True)
+        return tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True).strip()
+    except Exception as e:
+        print(f"Error during model generation for role {role}: {e}")
+        return f"error({e})"