quantum_head.py

"""
Q-GPT: Quantum-Enhanced GPT with Confidence Estimation
A quantum neural network head that estimates response confidence.

Author: squ11z1
"""

import torch
import torch.nn as nn
import numpy as np

try:
    import pennylane as qml
    PENNYLANE_AVAILABLE = True
except ImportError:
    PENNYLANE_AVAILABLE = False
    print("Warning: PennyLane not installed. Using classical fallback.")


class QuantumCircuit:
    """Variational Quantum Circuit for confidence estimation."""
    
    def __init__(self, n_qubits: int = 4, n_layers: int = 3):
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        
        if PENNYLANE_AVAILABLE:
            self.dev = qml.device("default.qubit", wires=n_qubits)
            self.circuit = qml.QNode(self._quantum_circuit, self.dev, interface="torch")
        
    def _quantum_circuit(self, inputs, weights):
        """
        Variational quantum circuit.
        
        Args:
            inputs: Input features [n_qubits]
            weights: Trainable parameters [n_layers, n_qubits, 3]
        """
        # Encode classical data into quantum states
        for i in range(self.n_qubits):
            qml.RY(inputs[i], wires=i)
            qml.RZ(inputs[i], wires=i)
        
        # Variational layers
        for layer in range(self.n_layers):
            # Rotation gates
            for i in range(self.n_qubits):
                qml.Rot(weights[layer, i, 0], 
                       weights[layer, i, 1], 
                       weights[layer, i, 2], wires=i)
            
            # Entanglement (CNOT ladder)
            for i in range(self.n_qubits - 1):
                qml.CNOT(wires=[i, i + 1])
            
            # Circular entanglement
            if self.n_qubits > 2:
                qml.CNOT(wires=[self.n_qubits - 1, 0])
        
        # Measure expectation values
        return [qml.expval(qml.PauliZ(i)) for i in range(self.n_qubits)]
    
    def forward(self, inputs, weights):
        """Execute quantum circuit."""
        if PENNYLANE_AVAILABLE:
            return self.circuit(inputs, weights)
        else:
            # Classical fallback: simple tanh transformation
            return torch.tanh(inputs @ weights.mean(dim=(0, 2)))


class QuantumHead(nn.Module):
    """
    Quantum-enhanced confidence estimation head for GPT.
    
    Takes hidden states from the last layer and outputs:
    - confidence: Estimated confidence in the response [0, 1]
    - uncertainty: Quantum-derived uncertainty measure
    """
    
    def __init__(
        self,
        hidden_size: int = 2880,  # GPT-OSS hidden size
        n_qubits: int = 4,
        n_layers: int = 3,
        intermediate_size: int = 64,
    ):
        super().__init__()
        
        self.hidden_size = hidden_size
        self.n_qubits = n_qubits
        self.n_layers = n_layers
        
        # Classical preprocessing: compress hidden states
        self.pre_quantum = nn.Sequential(
            nn.Linear(hidden_size, intermediate_size),
            nn.LayerNorm(intermediate_size),
            nn.GELU(),
            nn.Linear(intermediate_size, n_qubits),
            nn.Tanh(),  # Normalize to [-1, 1] for quantum encoding
        )
        
        # Quantum circuit
        self.quantum = QuantumCircuit(n_qubits, n_layers)
        
        # Quantum weights (trainable)
        self.quantum_weights = nn.Parameter(
            torch.randn(n_layers, n_qubits, 3) * 0.1
        )
        
        # Post-quantum processing
        self.post_quantum = nn.Sequential(
            nn.Linear(n_qubits, intermediate_size),
            nn.GELU(),
            nn.Linear(intermediate_size, 2),  # [confidence, uncertainty]
        )
        
        # Output heads
        self.confidence_activation = nn.Sigmoid()
        self.uncertainty_activation = nn.Softplus()
        
    def forward(self, hidden_states: torch.Tensor) -> dict:
        """
        Compute confidence and uncertainty from hidden states.
        
        Args:
            hidden_states: Last layer hidden states [batch, seq_len, hidden_size]
                          or pooled representation [batch, hidden_size]
        
        Returns:
            dict with 'confidence' and 'uncertainty' tensors
        """
        # Pool if sequence dimension exists
        if hidden_states.dim() == 3:
            # Use last token representation
            hidden_states = hidden_states[:, -1, :]
        
        batch_size = hidden_states.size(0)
        
        # Preprocess
        quantum_input = self.pre_quantum(hidden_states)  # [batch, n_qubits]
        
        # Process through quantum circuit (per sample)
        quantum_outputs = []
        for i in range(batch_size):
            qout = self.quantum.forward(
                quantum_input[i], 
                self.quantum_weights
            )
            if isinstance(qout, list):
                qout = torch.stack(qout)
            quantum_outputs.append(qout)
        
        quantum_output = torch.stack(quantum_outputs)  # [batch, n_qubits]
        
        # Post-process
        output = self.post_quantum(quantum_output)
        
        confidence = self.confidence_activation(output[:, 0])
        uncertainty = self.uncertainty_activation(output[:, 1])
        
        return {
            "confidence": confidence,
            "uncertainty": uncertainty,
            "should_refuse": confidence < 0.3,  # Low confidence = should refuse
        }
    
    def get_interpretable_confidence(self, confidence: torch.Tensor) -> str:
        """Convert confidence score to human-readable label."""
        conf = confidence.item() if confidence.dim() == 0 else confidence.mean().item()
        
        if conf >= 0.9:
            return "very high"
        elif conf >= 0.7:
            return "high"
        elif conf >= 0.5:
            return "moderate"
        elif conf >= 0.3:
            return "low"
        else:
            return "very low (consider refusing)"


class QGPT(nn.Module):
    """
    Q-GPT: GPT with Quantum Confidence Head
    
    Wraps any HuggingFace GPT model and adds quantum confidence estimation.
    """
    
    def __init__(self, base_model, quantum_head: QuantumHead = None):
        super().__init__()
        self.base_model = base_model
        
        # Get hidden size from model config
        if hasattr(base_model.config, 'hidden_size'):
            hidden_size = base_model.config.hidden_size
        elif hasattr(base_model.config, 'd_model'):
            hidden_size = base_model.config.d_model
        else:
            hidden_size = 2880  # GPT-OSS default
        
        self.quantum_head = quantum_head or QuantumHead(hidden_size=hidden_size)
        
    def forward(self, input_ids, attention_mask=None, **kwargs):
        """Forward pass with confidence estimation."""
        # Get base model outputs with hidden states
        outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            **kwargs
        )
        
        # Get last layer hidden states
        hidden_states = outputs.hidden_states[-1]
        
        # Compute quantum confidence
        confidence_output = self.quantum_head(hidden_states)
        
        # Add to outputs
        outputs.confidence = confidence_output["confidence"]
        outputs.uncertainty = confidence_output["uncertainty"]
        outputs.should_refuse = confidence_output["should_refuse"]
        
        return outputs
    
    def generate_with_confidence(
        self, 
        input_ids, 
        attention_mask=None,
        max_new_tokens=256,
        **kwargs
    ):
        """Generate text and return confidence score."""
        # Generate
        outputs = self.base_model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            output_hidden_states=True,
            return_dict_in_generate=True,
            **kwargs
        )
        
        # Get hidden states from last generation step
        if hasattr(outputs, 'hidden_states') and outputs.hidden_states:
            last_hidden = outputs.hidden_states[-1][-1]  # Last layer, last step
        else:
            # Fallback: run forward pass on generated sequence
            with torch.no_grad():
                model_outputs = self.base_model(
                    outputs.sequences,
                    output_hidden_states=True
                )
                last_hidden = model_outputs.hidden_states[-1]
        
        # Compute confidence
        confidence_output = self.quantum_head(last_hidden)
        
        return {
            "sequences": outputs.sequences,
            "confidence": confidence_output["confidence"],
            "uncertainty": confidence_output["uncertainty"],
            "should_refuse": confidence_output["should_refuse"],
            "confidence_label": self.quantum_head.get_interpretable_confidence(
                confidence_output["confidence"]
            ),
        }


def load_qgpt(
    model_name: str = "squ11z1/gpt-oss-9b-reasoning",
    quantum_head_path: str = None,
    device: str = "auto",
    torch_dtype = None,
    **kwargs
):
    """
    Load Q-GPT model with quantum head.
    
    Args:
        model_name: HuggingFace model name or path
        quantum_head_path: Path to trained quantum head weights
        device: Device to load model on
        torch_dtype: Model dtype (e.g., torch.bfloat16)
    
    Returns:
        QGPT model and tokenizer
    """
    from transformers import AutoModelForCausalLM, AutoTokenizer
    
    if torch_dtype is None:
        torch_dtype = torch.bfloat16
    
    # Load base model
    base_model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch_dtype,
        device_map=device,
        trust_remote_code=True,
        **kwargs
    )
    
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True,
        **kwargs
    )
    
    # Create Q-GPT
    model = QGPT(base_model)
    
    # Load quantum head weights if provided
    if quantum_head_path:
        state_dict = torch.load(quantum_head_path, map_location="cpu")
        model.quantum_head.load_state_dict(state_dict)
        print(f"Loaded quantum head from {quantum_head_path}")
    
    return model, tokenizer


if __name__ == "__main__":
    # Quick test
    print("Testing QuantumHead...")
    
    head = QuantumHead(hidden_size=2880)
    dummy_input = torch.randn(2, 2880)  # Batch of 2
    
    output = head(dummy_input)
    print(f"Confidence: {output['confidence']}")
    print(f"Uncertainty: {output['uncertainty']}")
    print(f"Should refuse: {output['should_refuse']}")
    
    print("\n✓ QuantumHead test passed!")