Build Your Own Jev: Open-Source System 1 Decision Engines

# build_your_own_jev.py
# A 30-line System 1 Decision Primitive using HuggingFace & PyTorch
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch.nn.functional as F

class LocalJevRouter:
    def __init__(self, model_id: str = "Qwen/Qwen2.5-1.5B-Instruct"):
        print(f"Loading local System 1 decision engine: {model_id}...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id, 
            torch_dtype=torch.float16, 
            device_map="auto"
        )

    def route(self, prompt: str, candidates: list[str]) -> dict:
        # Format single prompt framing options as explicit categorical letters
        letters = ["A", "B", "C", "D", "E"][:len(candidates)]
        options_text = "\n".join([f"{letter}: {opt}" for letter, opt in zip(letters, candidates)])
        input_text = f"Analyze task and select best category letter.\nTask: {prompt}\nOptions:\n{options_text}\nSelected Letter:"
        
        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model.device)
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            # Inspect next-token logits from the final forward pass
            next_token_logits = outputs.logits[0, -1, :]
            
            # Extract logits strictly for candidate letters
            candidate_token_ids = [self.tokenizer.encode(letter)[-1] for letter in letters]
            option_logits = next_token_logits[candidate_token_ids]
            
            # Softmax across candidate options to obtain calibrated probability distribution
            probabilities = F.softmax(option_logits, dim=-1).cpu().tolist()
            
        best_idx = int(torch.argmax(option_logits))
        return {
            "choice": candidates[best_idx],
            "confidence": round(probabilities[best_idx], 4),
            "distribution": {cand: round(prob, 4) for cand, prob in zip(candidates, probabilities)}
        }

# Usage Example:
if __name__ == "__main__":
    router = LocalJevRouter()
    result = router.route(
        prompt="Review pull request #402 for memory leak vulnerabilities",
        candidates=["code_search", "test_runner", "docs_lookup", "code_review", "none"]
    )
    print("Decision Output:", result)
    # Output: {'choice': 'code_review', 'confidence': 0.9642, ...}