1qwsd/LRM

1# [Use the ethics framework from ethical-rag-starter.py]2# Or minimal version:3 4from dataclasses import dataclass5from typing import List, Dict6from transformers import AutoTokenizer, AutoModelForCausalLM7import torch8 9@dataclass10class EthicsCheckResult:11    passed: bool12    score: float13    reasoning: str14    recommendations: List[str]15 16class AIEthicsFramework:17    BLOCKED_DOMAINS = ['medical_diagnosis_unsupervised', 'legal_judgment', 'hiring_decisions']18    19    def __init__(self):20        self.audit_log = []21    22    def validate_query(self, query: str) -> Dict:23        """Check if query is ethically acceptable"""24        pii_keywords = ['ssn', 'password', 'credit card']25        unsafe_words = ['hack', 'exploit', 'weaponize']26        27        has_pii = any(kw in query.lower() for kw in pii_keywords)28        is_unsafe = any(w in query.lower() for w in unsafe_words)29        30        is_allowed = not (has_pii or is_unsafe)31        reason = ""32        if has_pii:33            reason = "Query requests PII"34        elif is_unsafe:35            reason = "Query seeks harmful information"36        37        return {38            'is_allowed': is_allowed,39            'reason': reason or 'Query approved',40            'details': {'pii_check': has_pii, 'safety_check': is_unsafe}41        }42    43    def validate_response(self, response: str) -> EthicsCheckResult:44        """Validate generated response"""45        quality = len(response.split()) / 20  # Simple quality metric46        quality = min(quality, 1.0)47        48        return EthicsCheckResult(49            passed=quality > 0.3,50            score=quality,51            reasoning="Response quality acceptable" if quality > 0.3 else "Response too brief",52            recommendations=[]53        )54 55def initialize_llm(model_name: str):56    """Load and initialize LLM"""57    tokenizer = AutoTokenizer.from_pretrained(model_name)58    model = AutoModelForCausalLM.from_pretrained(59        model_name,60        torch_dtype=torch.float16,61        device_map="auto",62        load_in_8bit=True  # For memory efficiency63    )64    65    class SimpleLLM:66        def __init__(self, model, tokenizer):67            self.model = model68            self.tokenizer = tokenizer69        70        def generate(self, prompt: str, max_tokens: int = 300):71            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)72            with torch.no_grad():73                outputs = self.model.generate(74                    **inputs,75                    max_new_tokens=max_tokens,76                    temperature=0.7,77                    top_p=0.978                )79            return self.tokenizer.decode(outputs, skip_special_tokens=True)80    81    return SimpleLLM(model, tokenizer)82