eval

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2Model
from datasets import load_dataset
from main import ConfidenceEnhancedTransformer  # Import the class from main.py

# Load the trained model and tokenizer
model_name = 'confidence_model'
model_path = model_name
tokenizer_path = model_name
tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_path)

model = ConfidenceEnhancedTransformer.from_pretrained(model_path, attn_implementation="eager")

# Move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Set the model to evaluation mode
model.eval()

# Load the WikiText-2 dataset
dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')

# Ensure the example text is not empty and preprocess it
example_text = dataset['text'][3].strip()  # Get the first example from the dataset and strip whitespace
if not example_text:
    raise ValueError("The example text from the dataset is empty. Please check the dataset.")

# Example usage
if __name__ == "__main__":
    prompt = "tefewafwef aoasdfsfasdfsadfdfasdsdijfoiwej"
    inputs = tokenizer(prompt, return_tensors='pt').to(device)

    # Check if input_ids is empty
    if inputs.input_ids.size(1) == 0:
        raise ValueError("The input text resulted in an empty input_ids tensor. Please check the input text.")

    with torch.no_grad():
        outputs = model(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            num_dropout_samples=10  # Increased from 5 to 10
        )
    # Get the confidence score
    confidence_score = outputs['confidence_score'].item()
    ood_score = outputs['ood_score'].item()
    print(f"Refined Confidence Score: {confidence_score}")
    print(f"OOD Score: {ood_score}")

    # Decode and print the generated text (not part of the confidence mechanism)
    generated_text = tokenizer.decode(outputs['lm_logits'].argmax(-1).squeeze().tolist())
    print(f"OOD example: {prompt}")
    print(f"Generated Text: {generated_text}")

    # Evaluate on an in-distribution example from WikiText-2
    inputs = tokenizer(example_text, return_tensors='pt').to(device)

    # Check if input_ids is empty
    if inputs.input_ids.size(1) == 0:
        raise ValueError("The example text resulted in an empty input_ids tensor. Please check the example text.")

    with torch.no_grad():
        outputs = model(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            num_dropout_samples=10
        )
    # Get the confidence score for the in-distribution example
    confidence_score = outputs['confidence_score'].item()
    ood_score = outputs['ood_score'].item()
    print(f"In-Distribution Example - Refined Confidence Score: {confidence_score}")
    print(f"In-Distribution Example - OOD Score: {ood_score}")

    # Decode and print the generated text for the in-distribution example
    generated_text = tokenizer.decode(outputs['lm_logits'].argmax(-1).squeeze().tolist())
    print(f"example text: {example_text}")
    print(f"In-Distribution Example - Generated Text: {generated_text}")


main

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Model, GPT2Tokenizer
import torch.nn.functional as F

class ConfidenceEnhancedTransformer(GPT2LMHeadModel):
    def __init__(self, config):
        super(ConfidenceEnhancedTransformer, self).__init__(config)
        self.transformer = GPT2Model(config)
        #self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)  # Language modeling head

        # Confidence scoring head for epistemic uncertainty and OOD detection
        self.confidence_head = nn.Sequential(
            nn.Linear(config.n_embd, 128),
            nn.ReLU(),
            nn.Linear(128, 1),  # Single output for confidence score
            nn.Sigmoid()  # Confidence score between 0 and 1
        )

        # OOD detector head
        self.ood_detector = nn.Linear(config.n_embd, 1)  # Auxiliary head for OOD detection
        self.init_weights()

    def forward(self, input_ids, attention_mask=None, labels=None, num_dropout_samples=10):
        # Standard forward pass through transformer
        outputs = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            output_attentions=True,
            output_hidden_states=True,
            return_dict=True,
        )
        lm_logits = outputs.logits  # [batch_size, sequence_length, vocab_size]
        hidden_states = outputs.hidden_states[-1]  # Get the last hidden state
        attentions = outputs.attentions  # Attention weights

        # Base confidence score from hidden states
        base_confidence_score = self.confidence_head(hidden_states.mean(dim=1))

        # Attention-based confidence signal
        attention_entropy = []
        for attn_layer in attentions:
            attn_probs = attn_layer.mean(dim=1)  # Mean over heads
            attn_entropy = -torch.sum(attn_probs * torch.log(attn_probs + 1e-12), dim=-1)
            attention_entropy.append(attn_entropy.mean(dim=-1))  # Mean over tokens
        avg_attention_entropy = torch.stack(attention_entropy).mean(dim=0)  # Mean over layers

        # Monte Carlo Dropout for variance estimation
        variance_confidence = 0.0
        dropout_scores = []
        if num_dropout_samples > 1:
            original_mode = self.training  # Save original mode
            self.train()  # Enable dropout layers
            for _ in range(num_dropout_samples):
                # Removed torch.no_grad() to ensure dropout behaves correctly
                dropout_outputs = super().forward(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=True,
                    return_dict=True,
                )
                dropout_hidden_states = dropout_outputs.hidden_states[-1]
                dropout_confidence = self.confidence_head(dropout_hidden_states.mean(dim=1))
                dropout_scores.append(dropout_confidence)

            self.train(original_mode)  # Restore original mode
            # Calculate variance of dropout predictions as a confidence measure
            dropout_scores = torch.stack(dropout_scores)  # [num_samples, batch_size, 1]
            variance_confidence = torch.var(dropout_scores, dim=0).mean()
        else:
            variance_confidence = torch.tensor(0.0).to(hidden_states.device)

        # OOD detection score
        ood_score = torch.sigmoid(self.ood_detector(hidden_states.mean(dim=1))).squeeze()

        # Adjust this line to handle ood_score shape correctly
        if len(ood_score.shape) == 0:
            ood_score = ood_score.unsqueeze(0)  # Add batch dimension if it's a scalar

        # Combine all signals into a refined confidence score
        refined_confidence_score = (
            base_confidence_score
            - variance_confidence  # Lower confidence if high variance
            - avg_attention_entropy.unsqueeze(1)  # Lower confidence if high attention entropy
            - ood_score.unsqueeze(1)  # Lower confidence if high OOD score
        ).clamp(0, 1)  # Ensure the final score is between 0 and 1

        # Calculate total loss if labels are provided
        total_loss = None
        if labels is not None:
            # Language modeling loss
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            lm_loss = nn.CrossEntropyLoss()(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

            # Entropy-based confidence loss
            token_probs = F.softmax(shift_logits, dim=-1)
            entropy = -torch.sum(token_probs * torch.log(token_probs + 1e-12), dim=-1).mean(dim=-1)
            confidence_loss = nn.MSELoss()(base_confidence_score.squeeze(), 1 - entropy)

            # OOD loss
            ood_loss = nn.BCELoss()(ood_score, torch.zeros_like(ood_score))  # Penalty for high OOD score

            # Combine all losses with adjusted weights
            total_loss = lm_loss + 0.5 * confidence_loss + 0.3 * ood_loss  # Increased weights

        return {
            'loss': total_loss,
            'lm_logits': lm_logits,
            'confidence_score': refined_confidence_score,  # Refined confidence score
            'ood_score': ood_score,  # Out-of-distribution score
            'base_confidence_score': base_confidence_score.squeeze(),  # Added for logging
            'variance_confidence': variance_confidence,  # Added for logging
            'avg_attention_entropy': avg_attention_entropy  # Added for logging
        }

# Example usage of the model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = ConfidenceEnhancedTransformer.from_pretrained('gpt2', attn_implementation='eager')

# Example input
input_text = "What is the capital of USA?"
input_tokens = tokenizer(input_text, return_tensors="pt")
outputs = model(input_tokens['input_ids'], num_dropout_samples=10)

# Get the confidence score
confidence_score = outputs['confidence_score'].item()
ood_score = outputs['ood_score'].item()
print(f"Refined Confidence Score: {confidence_score}")
print(f"OOD Score: {ood_score}")

# Decode and print the generated text (not part of the confidence mechanism)
generated_text = tokenizer.decode(outputs['lm_logits'].argmax(-1).squeeze().tolist())
print(f"Generated Text: {generated_text}")


import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Model, GPT2Tokenizer
import torch.nn.functional as F

class ConfidenceEnhancedTransformer(GPT2LMHeadModel):
    def __init__(self, config):
        super(ConfidenceEnhancedTransformer, self).__init__(config)
        self.transformer = GPT2Model(config)
        #self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)  # Language modeling head

        # Confidence scoring head for epistemic uncertainty and OOD detection
        self.confidence_head = nn.Sequential(
            nn.Linear(config.n_embd, 128),
            nn.ReLU(),
            nn.Linear(128, 1),  # Single output for confidence score
            nn.Sigmoid()  # Confidence score between 0 and 1
        )

        # OOD detector head
        self.ood_detector = nn.Linear(config.n_embd, 1)  # Auxiliary head for OOD detection
        self.init_weights()

    def forward(self, input_ids, attention_mask=None, labels=None, num_dropout_samples=10):
        # Standard forward pass through transformer
        outputs = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            output_attentions=True,
            output_hidden_states=True,
            return_dict=True,
        )
        lm_logits = outputs.logits  # [batch_size, sequence_length, vocab_size]
        hidden_states = outputs.hidden_states[-1]  # Get the last hidden state
        attentions = outputs.attentions  # Attention weights

        # Base confidence score from hidden states
        base_confidence_score = self.confidence_head(hidden_states.mean(dim=1))

        # Attention-based confidence signal
        attention_entropy = []


forward

import torch.nn.functional as F

class ConfidenceEnhancedTransformer(GPT2LMHeadModel):
    def __init__(self, config):
        super(ConfidenceEnhancedTransformer, self).__init__(config)
        self.transformer = GPT2Model(config)
        #self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)  # Language modeling head

        # Confidence scoring head for epistemic uncertainty and OOD detection
        self.confidence_head = nn.Sequential(
            nn.Linear(config.n_embd, 128),
            nn.ReLU(),
            nn.Linear(128, 1),  # Single output for confidence score
            nn.Sigmoid()  # Confidence score between 0 and 1
        )

        # OOD detector head
        self.ood_detector = nn.Linear(config.n_embd, 1)  # Auxiliary head for OOD detection
        self.init_weights()

    def forward(self, input_ids, attention_mask=None, labels=None, num_dropout_samples=10):
        # Standard forward pass through transformer
        outputs = super().forward(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            output_attentions=True,
            output_hidden_states=True,
            return_dict=True,
        )
        lm_logits = outputs.logits  # [batch_size, sequence_length, vocab_size]
        hidden_states = outputs.hidden_states[-1]  # Get the last hidden state
        attentions = outputs.attentions  # Attention weights

        # Base confidence score from hidden states
        base_confidence_score = self.confidence_head(hidden_states.mean(dim=1))

        # Attention-based confidence signal
        attention_entropy = []
        for attn_layer in attentions:
            attn_probs = attn_layer.mean(dim=1)  # Mean over heads
            attn_entropy = -torch.sum(attn_probs * torch.log(attn_probs + 1e-12), dim=-1)
            attention_entropy.append(attn_entropy.mean(dim=-1))  # Mean over tokens
        avg_attention_entropy = torch.stack(attention_entropy).mean(dim=0)  # Mean over layers

        # Monte Carlo Dropout for variance estimation
        variance_confidence = 0.0
        dropout_scores = []
        if num_dropout_samples > 1:
            original_mode = self.training  # Save original mode
            self.train()  # Enable dropout layers
            for _ in range(num_dropout_samples):
                # Removed torch.no_grad() to ensure dropout behaves correctly
                dropout_outputs = super().forward(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=True,
                    return_dict=True,
                )
                dropout_hidden_states = dropout_outputs.hidden_states[-1]
                dropout_confidence = self.confidence_head(dropout_hidden_states.mean(dim=1))
                dropout_scores.append(dropout_confidence)

            self.train(original_mode)  # Restore original mode
            # Calculate variance of dropout predictions as a confidence measure
            dropout_scores = torch.stack(dropout_scores)  # [num_samples, batch_size, 1]
            variance_confidence = torch.var(dropout_scores, dim=0).mean()
        else:
            variance_confidence = torch.tensor(0.0).to(hidden_states.device)

        # OOD detection score
        ood_score = torch.sigmoid(self.ood_detector(hidden_states.mean(dim=1))).squeeze()

        # Adjust this line to handle ood_score shape correctly
        if len(ood_score.shape) == 0:
            ood_score = ood_score.unsqueeze(0)  # Add batch dimension if it's a scalar

        # Combine all signals into a refined confidence score
        refined_confidence_score = (
            base_confidence_score
            - variance_confidence  # Lower confidence if high variance
            - avg_attention_entropy.unsqueeze(1)  # Lower confidence if high attention entropy
            - ood_score.unsqueeze(1)  # Lower confidence if high OOD score
        ).clamp(0, 1)  # Ensure the final score is between 0 and 1

        # Calculate total loss if labels are provided
        total_loss = None
        if labels is not None:
            # Language modeling loss
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            lm_loss = nn.CrossEntropyLoss()(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

            # Entropy-based confidence loss
            token_probs = F.softmax(shift_logits, dim=-1)
            entropy = -torch.sum(token_probs * torch.log(token_probs + 1e-12), dim=-1).mean(dim=-1)
            confidence_loss = nn.MSELoss()(base_confidence_score.squeeze(), 1 - entropy)

            # OOD loss
            ood_loss = nn.BCELoss()(ood_score, torch.zeros_like(ood_score))  # Penalty for high OOD score

            # Combine all losses with adjusted weights
            total_loss = lm_loss + 0.5 * confidence_loss + 0.3 * ood_loss  # Increased weights

        return {
            'loss': total_loss,
            'lm_logits': lm_logits,
            'confidence_score': refined_confidence_score,  # Refined confidence score
            'ood_score': ood_score,  # Out-of-distribution score
            'base_confidence_score': base_confidence_score.squeeze(),  # Added for logging
            'variance_confidence': variance_confidence,  # Added for logging
            'avg_attention_entropy': avg_attention_entropy  # Added for logging
        }

# Example usage of the model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = ConfidenceEnhancedTransformer.from_pretrained('gpt2', attn_implementation='eager')

# Example input
input_text = "What is the capital of USA?"
input_tokens = tokenizer(input_text, return_tensors="pt")
outputs = model(input_tokens['input_ids'], num_dropout_samples=10)

# Get the confidence score
confidence_score = outputs['confidence_score'].item()
ood_score = outputs['ood_score'].item()
print(f"Refined Confidence Score: {confidence_score}")
print(f"OOD Score: {ood_score}")

# Decode and print the generated text (not part of the confidence mechanism)
generated_text = tokenizer.decode(outputs['lm_logits'].argmax(-1).squeeze().tolist())
print(f"Generated Text: {generated_text}")


train

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer, get_linear_schedule_with_warmup
from tqdm import tqdm
import torch.optim as optim
from datasets import load_dataset
from main import ConfidenceEnhancedTransformer
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import matplotlib.pyplot as plt

num_epochs = 3
save_steps = 500  # Save the model every 500 steps

# Define a custom dataset
class TextDataset(Dataset):
    def __init__(self, tokenizer, texts, block_size=128):
        self.examples = []

        for text in texts:
            tokenized_text = tokenizer.encode(text)
            for i in range(0, len(tokenized_text) - block_size + 1, block_size):
                self.examples.append(
                    torch.tensor(tokenized_text[i:i + block_size], dtype=torch.long)
                )

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return self.examples[i]

# Load tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')

# Load the WikiText-2 dataset
dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
texts = dataset['text']

# Split the dataset into training and validation sets
train_texts, val_texts = train_test_split(texts, test_size=0.1, random_state=42)

# Prepare datasets
train_dataset = TextDataset(
    tokenizer=tokenizer,
    texts=train_texts,
    block_size=128
)
val_dataset = TextDataset(
    tokenizer=tokenizer,
    texts=val_texts,
    block_size=128
)

# Create DataLoaders
train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=4, shuffle=False)

# Prepare optimizer and scheduler
optimizer = optim.AdamW(model.parameters(), lr=5e-5)
total_steps = len(train_dataloader) * num_epochs
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=0,
    num_training_steps=total_steps
)

# Move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Define ECE computation function
def compute_ece(preds, confidences, n_bins=10):
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    for i in range(n_bins):
        bin_lower = bin_boundaries[i]
        bin_upper = bin_boundaries[i + 1]
        in_bin = (confidences > bin_lower) & (confidences <= bin_upper)
        prop_in_bin = np.mean(in_bin)
        if prop_in_bin > 0:
            accuracy_in_bin = np.mean(preds[in_bin] == preds[in_bin])  # Adjust based on your prediction mechanism
            avg_confidence_in_bin = np.mean(confidences[in_bin])
            ece += np.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
    return ece

# Define Reliability Diagram function
def reliability_diagram(confidences, predictions, n_bins=10):
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    bin_centers = (bin_boundaries[:-1] + bin_boundaries[1:]) / 2
    accuracy = np.zeros(n_bins)
    confidence = np.zeros(n_bins)
    prop = np.zeros(n_bins)

    for i in range(n_bins):
        in_bin = (confidences > bin_boundaries[i]) & (confidences <= bin_boundaries[i + 1])
        prop[i] = np.mean(in_bin)
        if prop[i] > 0:
            accuracy[i] = np.mean(predictions[in_bin] == predictions[in_bin])  # Replace with actual labels
            confidence[i] = np.mean(confidences[in_bin])

    plt.figure(figsize=(8, 6))
    plt.plot(bin_centers, accuracy, marker='o', label='Accuracy')
    plt.plot(bin_centers, confidence, marker='s', label='Confidence')
    plt.fill_between(bin_boundaries[:-1], 0, 1, color='gray', alpha=0.1)
    plt.xlabel('Confidence')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.title('Reliability Diagram')
    plt.show()

# Training loop
model.train()
global_step = 0
for epoch in range(num_epochs):
    print(f"Epoch {epoch + 1}/{num_epochs}")
    epoch_loss = 0
    for batch in tqdm(train_dataloader):
        inputs = batch.to(device)
        labels = inputs.clone()

        optimizer.zero_grad()

        outputs = model(
            input_ids=inputs,
            labels=labels,
            num_dropout_samples=15  # You can adjust this number
        )
        loss = outputs['loss']
        loss.backward()

        # Check gradients for confidence_head and ood_detector
        for name, param in model.named_parameters():
            if 'confidence_head' in name or 'ood_detector' in name:
                if param.grad is not None:
                    print(f"Gradient for {name}: {param.grad.mean().item()}")

        optimizer.step()
        scheduler.step()

        epoch_loss += loss.item()
        global_step += 1

        # Log intermediate values for debugging
        base_confidence_score = outputs['base_confidence_score'].mean().item()
        variance_confidence = outputs['variance_confidence'].item()
        avg_attention_entropy = outputs['avg_attention_entropy'].mean().item()
        ood_score = outputs['ood_score'].mean().item()

        print(f"Step {global_step}: Loss = {loss.item():.4f}, "
              f"Base Confidence Score = {base_confidence_score:.4f}, "
              f"Variance Confidence = {variance_confidence:.4f}, "
              f"Avg Attention Entropy = {avg_attention_entropy:.4f}, "
              f"OOD Score = {ood_score:.4f}")

        # Save the model every 500 steps
        if global_step % save_steps == 0:
            model.save_pretrained(f'model_step_{global_step}.pth')
            tokenizer.save_pretrained(f'tokenizer_step_{global_step}.pth')

    avg_train_loss = epoch_loss / len(train_dataloader)
    print(f"Average Training Loss: {avg_train_loss:.4f}")

    # Validation phase
    model.eval()
    val_loss = 0
    all_val_preds = []
    all_val_confidences = []
    with torch.no_grad():
        for val_batch in tqdm(val_dataloader, desc="Validation"):
            val_inputs = val_batch.to(device)
            val_labels = val_inputs.clone()

            val_outputs = model(
                input_ids=val_inputs,
                labels=val_labels,
                num_dropout_samples=15
            )
            val_loss += val_outputs['loss'].item()

            # Collect predictions and confidence scores for calibration
            val_confidences = val_outputs['base_confidence_score'].cpu().numpy()
            # Assuming you're using the model's logits to derive predictions
            val_logits = val_outputs['logits']
            val_preds = torch.argmax(val_logits, dim=-1).cpu().numpy()
            all_val_preds.extend(val_preds.flatten())
            all_val_confidences.extend(val_confidences.flatten())

    avg_val_loss = val_loss / len(val_dataloader)
    print(f"Average Validation Loss after Epoch {epoch + 1}: {avg_val_loss:.4f}")

    # Calculate ECE
    ece = compute_ece(
        preds=np.array(all_val_preds),
        confidences=np.array(all_val_confidences),
        n_bins=10
    )
    print(f"Expected Calibration Error (ECE) after Epoch {epoch + 1}: {ece:.4f}")

    # Plot Reliability Diagram
    reliability_diagram(
        confidences=np.array(all_val_confidences),
        predictions=np.array(all_val_preds),
        n_bins=10
    )

    # Reset model to training mode
    model.train()

# Save the final trained model
model.save_pretrained('confidence_model')
tokenizer.save_pretrained('confidence_model')


TextDataset

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer, get_linear_schedule_with_warmup
from tqdm import tqdm
import torch.optim as optim
from datasets import load_dataset
from main import ConfidenceEnhancedTransformer
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import matplotlib.pyplot as plt

num_epochs = 3
save_steps = 500  # Save the model every 500 steps

# Define a custom dataset
class TextDataset(Dataset):
    def __init__(self, tokenizer, texts, block_size=128):
        self.examples = []

        for text in texts:
            tokenized_text = tokenizer.encode(text)
            for i in range(0, len(tokenized_text) - block_size + 1, block_size):
                self.examples.append(
                    torch.tensor(tokenized_text[i:i + block_size], dtype=torch.long)
                )

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return self.examples[i]

# Load tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')

# Load the WikiText-2 dataset
dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
texts = dataset['text']

# Split the dataset into training and validation sets
train_texts, val_texts = train_test_split(texts, test_size=0.1, random_state=42)

# Prepare datasets
train_dataset = TextDataset(
    tokenizer=tokenizer,
    texts=train_texts,
    block_size=128
)
val_dataset = TextDataset(
    tokenizer=tokenizer,


)

# Create DataLoaders
train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)
val_dataloader = DataLoader(val_dataset, batch_size=4, shuffle=False)

# Prepare optimizer and scheduler
optimizer = optim.AdamW(model.parameters(), lr=5e-5)
total_steps = len(train_dataloader) * num_epochs
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=0,
    num_training_steps=total_steps
)

# Move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Define ECE computation function
def compute_ece(preds, confidences, n_bins=10):
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    for i in range(n_bins):
        bin_lower = bin_boundaries[i]
        bin_upper = bin_boundaries[i + 1]
        in_bin = (confidences > bin_lower) & (confidences <= bin_upper)
        prop_in_bin = np.mean(in_bin)
        if prop_in_bin > 0:
            accuracy_in_bin = np.mean(preds[in_bin] == preds[in_bin])  # Adjust based on your prediction mechanism
            avg_confidence_in_bin = np.mean(confidences[in_bin])
            ece += np.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
    return ece

# Define Reliability Diagram function
def reliability_diagram(confidences, predictions, n_bins=10):
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    bin_centers = (bin_boundaries[:-1] + bin_boundaries[1:]) / 2
    accuracy = np.zeros(n_bins)
    confidence = np.zeros(n_bins)
    prop = np.zeros(n_bins)

    for i in range(n_bins):
        in_bin = (confidences > bin_boundaries[i]) & (confidences <= bin_boundaries[i + 1])
        prop[i] = np.mean(in_bin)
        if prop[i] > 0:
            accuracy[i] = np.mean(predictions[in_bin] == predictions[in_bin])  # Replace with actual labels
            confidence[i] = np.mean(confidences[in_bin])

    plt.figure(figsize=(8, 6))
    plt.plot(bin_centers, accuracy, marker='o', label='Accuracy')
    plt.plot(bin_centers, confidence, marker='s', label='Confidence')
    plt.fill_between(bin_boundaries[:-1], 0, 1, color='gray', alpha=0.1)


# Move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Define ECE computation function
def compute_ece(preds, confidences, n_bins=10):
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    for i in range(n_bins):
        bin_lower = bin_boundaries[i]
        bin_upper = bin_boundaries[i + 1]
        in_bin = (confidences > bin_lower) & (confidences <= bin_upper)
        prop_in_bin = np.mean(in_bin)
        if prop_in_bin > 0:
            accuracy_in_bin = np.mean(preds[in_bin] == preds[in_bin])  # Adjust based on your prediction mechanism
            avg_confidence_in_bin = np.mean(confidences[in_bin])
            ece += np.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
    return ece

# Define Reliability Diagram function
def reliability_diagram(confidences, predictions, n_bins=10):
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    bin_centers = (bin_boundaries[:-1] + bin_boundaries[1:]) / 2
    accuracy = np.zeros(n_bins)
    confidence = np.zeros(n_bins)
    prop = np.zeros(n_bins)

    for i in range(n_bins):
        in_bin = (confidences > bin_boundaries[i]) & (confidences <= bin_boundaries[i + 1])
        prop[i] = np.mean(in_bin)
        if prop[i] > 0:
            accuracy[i] = np.mean(predictions[in_bin] == predictions[in_bin])  # Replace with actual labels
            confidence[i] = np.mean(confidences[in_bin])

    plt.figure(figsize=(8, 6))
    plt.plot(bin_centers, accuracy, marker='o', label='Accuracy')
    plt.plot(bin_centers, confidence, marker='s', label='Confidence')
    plt.fill_between(bin_boundaries[:-1], 0, 1, color='gray', alpha=0.1)
    plt.xlabel('Confidence')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.title('Reliability Diagram')
    plt.show()

# Training loop
model.train()
global_step = 0
for epoch in range(num_epochs):
    print(f"Epoch {epoch + 1}/{num_epochs}")
    epoch_loss = 0
    for batch in tqdm(train_dataloader):
        inputs = batch.to(device)
        labels = inputs.clone()

        optimizer.zero_grad()

        outputs = model(
            input_ids=inputs,
            labels=labels,
            num_dropout_samples=15  # You can adjust this number
        )
        loss = outputs['loss']
        loss.backward()


import torch
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer, get_linear_schedule_with_warmup
from tqdm import tqdm
import torch.optim as optim
from datasets import load_dataset
from main import ConfidenceEnhancedTransformer
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import matplotlib.pyplot as plt

num_epochs = 3
save_steps = 500  # Save the model every 500 steps

# Define a custom dataset
class TextDataset(Dataset):
    def __init__(self, tokenizer, texts, block_size=128):
        self.examples = []

        for text in texts:
            tokenized_text = tokenizer.encode(text)
            for i in range(0, len(tokenized_text) - block_size + 1, block_size):
                self.examples.append(
                    torch.tensor(tokenized_text[i:i + block_size], dtype=torch.long)
                )

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return self.examples[i]

# Load tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')

# Load the WikiText-2 dataset
dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
texts = dataset['text']

# Split the dataset into training and validation sets
train_texts, val_texts = train_test_split(texts, test_size=0.1, random_state=42)

# Prepare datasets
train_dataset = TextDataset(


from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import matplotlib.pyplot as plt

num_epochs = 3
save_steps = 500  # Save the model every 500 steps

# Define a custom dataset
class TextDataset(Dataset):
    def __init__(self, tokenizer, texts, block_size=128):
        self.examples = []

        for text in texts:
            tokenized_text = tokenizer.encode(text)
            for i in range(0, len(tokenized_text) - block_size + 1, block_size):
                self.examples.append(
                    torch.tensor(tokenized_text[i:i + block_size], dtype=torch.long)
                )

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return self.examples[i]

# Load tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')

# Load the WikiText-2 dataset
dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
texts = dataset['text']

# Split the dataset into training and validation sets
train_texts, val_texts = train_test_split(texts, test_size=0.1, random_state=42)

# Prepare datasets
train_dataset = TextDataset(
    tokenizer=tokenizer,
    texts=train_texts,
    block_size=128


import matplotlib.pyplot as plt

num_epochs = 3
save_steps = 500  # Save the model every 500 steps

# Define a custom dataset
class TextDataset(Dataset):
    def __init__(self, tokenizer, texts, block_size=128):
        self.examples = []

        for text in texts:
            tokenized_text = tokenizer.encode(text)
            for i in range(0, len(tokenized_text) - block_size + 1, block_size):
                self.examples.append(
                    torch.tensor(tokenized_text[i:i + block_size], dtype=torch.long)
                )

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return self.examples[i]

# Load tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')

# Load the WikiText-2 dataset
dataset = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
texts = dataset['text']

# Split the dataset into training and validation sets
train_texts, val_texts = train_test_split(texts, test_size=0.1, random_state=42)

# Prepare datasets
train_dataset = TextDataset(
    tokenizer=tokenizer,
    texts=train_texts,
    block_size=128
)
val_dataset = TextDataset(
    tokenizer=tokenizer,


test

import unittest
import torch
from transformers import GPT2Tokenizer
from main import ConfidenceEnhancedTransformer

class TestConfidenceEnhancedTransformer(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        cls.model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')
        cls.model.eval()  # Set model to evaluation mode

    def test_initialization(self):
        self.assertIsInstance(self.model, ConfidenceEnhancedTransformer)

    def test_forward_pass(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        self.assertIn('lm_logits', outputs)
        self.assertIn('confidence_score', outputs)
        self.assertIn('ood_score', outputs)

    def test_confidence_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        confidence_score = outputs['confidence_score'].item()
        self.assertGreaterEqual(confidence_score, 0.0)
        self.assertLessEqual(confidence_score, 1.0)

    def test_ood_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        ood_score = outputs['ood_score'].item()
        self.assertGreaterEqual(ood_score, 0.0)
        self.assertLessEqual(ood_score, 1.0)

    def test_loss_calculation(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)

        self.assertIn('loss', outputs)
        self.assertIsNotNone(outputs['loss'])

    def test_training_step(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=5e-5)

        self.model.train()
        optimizer.zero_grad()
        outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)
        loss = outputs['loss']
        loss.backward()
        optimizer.step()

        self.assertIsNotNone(loss)
        self.assertGreater(loss.item(), 0.0)

if __name__ == '__main__':
    unittest.main()


setUpClass

import unittest
import torch
from transformers import GPT2Tokenizer
from main import ConfidenceEnhancedTransformer

class TestConfidenceEnhancedTransformer(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        cls.model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')
        cls.model.eval()  # Set model to evaluation mode

    def test_initialization(self):
        self.assertIsInstance(self.model, ConfidenceEnhancedTransformer)

    def test_forward_pass(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        self.assertIn('lm_logits', outputs)
        self.assertIn('confidence_score', outputs)
        self.assertIn('ood_score', outputs)

    def test_confidence_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)


import unittest
import torch
from transformers import GPT2Tokenizer
from main import ConfidenceEnhancedTransformer

class TestConfidenceEnhancedTransformer(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        cls.model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')
        cls.model.eval()  # Set model to evaluation mode

    def test_initialization(self):
        self.assertIsInstance(self.model, ConfidenceEnhancedTransformer)

    def test_forward_pass(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        self.assertIn('lm_logits', outputs)
        self.assertIn('confidence_score', outputs)
        self.assertIn('ood_score', outputs)

    def test_confidence_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        confidence_score = outputs['confidence_score'].item()
        self.assertGreaterEqual(confidence_score, 0.0)
        self.assertLessEqual(confidence_score, 1.0)


import unittest
import torch
from transformers import GPT2Tokenizer
from main import ConfidenceEnhancedTransformer

class TestConfidenceEnhancedTransformer(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        cls.model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')
        cls.model.eval()  # Set model to evaluation mode

    def test_initialization(self):
        self.assertIsInstance(self.model, ConfidenceEnhancedTransformer)

    def test_forward_pass(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        self.assertIn('lm_logits', outputs)
        self.assertIn('confidence_score', outputs)
        self.assertIn('ood_score', outputs)

    def test_confidence_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        confidence_score = outputs['confidence_score'].item()
        self.assertGreaterEqual(confidence_score, 0.0)
        self.assertLessEqual(confidence_score, 1.0)

    def test_ood_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        ood_score = outputs['ood_score'].item()
        self.assertGreaterEqual(ood_score, 0.0)
        self.assertLessEqual(ood_score, 1.0)


    @classmethod
    def setUpClass(cls):
        cls.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        cls.model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')
        cls.model.eval()  # Set model to evaluation mode

    def test_initialization(self):
        self.assertIsInstance(self.model, ConfidenceEnhancedTransformer)

    def test_forward_pass(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        self.assertIn('lm_logits', outputs)
        self.assertIn('confidence_score', outputs)
        self.assertIn('ood_score', outputs)

    def test_confidence_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        confidence_score = outputs['confidence_score'].item()
        self.assertGreaterEqual(confidence_score, 0.0)
        self.assertLessEqual(confidence_score, 1.0)

    def test_ood_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        ood_score = outputs['ood_score'].item()
        self.assertGreaterEqual(ood_score, 0.0)
        self.assertLessEqual(ood_score, 1.0)

    def test_loss_calculation(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)

        self.assertIn('loss', outputs)
        self.assertIsNotNone(outputs['loss'])


    def test_forward_pass(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        self.assertIn('lm_logits', outputs)
        self.assertIn('confidence_score', outputs)
        self.assertIn('ood_score', outputs)

    def test_confidence_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        confidence_score = outputs['confidence_score'].item()
        self.assertGreaterEqual(confidence_score, 0.0)
        self.assertLessEqual(confidence_score, 1.0)

    def test_ood_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        ood_score = outputs['ood_score'].item()
        self.assertGreaterEqual(ood_score, 0.0)
        self.assertLessEqual(ood_score, 1.0)

    def test_loss_calculation(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)

        self.assertIn('loss', outputs)
        self.assertIsNotNone(outputs['loss'])

    def test_training_step(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=5e-5)

        self.model.train()
        optimizer.zero_grad()
        outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)


    def test_confidence_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        confidence_score = outputs['confidence_score'].item()
        self.assertGreaterEqual(confidence_score, 0.0)
        self.assertLessEqual(confidence_score, 1.0)

    def test_ood_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        ood_score = outputs['ood_score'].item()
        self.assertGreaterEqual(ood_score, 0.0)
        self.assertLessEqual(ood_score, 1.0)

    def test_loss_calculation(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)

        self.assertIn('loss', outputs)
        self.assertIsNotNone(outputs['loss'])

    def test_training_step(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=5e-5)

        self.model.train()
        optimizer.zero_grad()
        outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)
        loss = outputs['loss']
        loss.backward()
        optimizer.step()

        self.assertIsNotNone(loss)
        self.assertGreater(loss.item(), 0.0)

if __name__ == '__main__':
    unittest.main()


    def test_ood_score_range(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)

        ood_score = outputs['ood_score'].item()
        self.assertGreaterEqual(ood_score, 0.0)
        self.assertLessEqual(ood_score, 1.0)

    def test_loss_calculation(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)

        self.assertIn('loss', outputs)
        self.assertIsNotNone(outputs['loss'])

    def test_training_step(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        labels = input_tokens['input_ids'].clone()
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=5e-5)

        self.model.train()
        optimizer.zero_grad()
        outputs = self.model(input_tokens['input_ids'], labels=labels, num_dropout_samples=2)
        loss = outputs['loss']
        loss.backward()
        optimizer.step()

        self.assertIsNotNone(loss)
        self.assertGreater(loss.item(), 0.0)

if __name__ == '__main__':
    unittest.main()


eval.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPT2Model
from datasets import load_dataset
from main import ConfidenceEnhancedTransformer  # Import the class from main.py

# Load the trained model and tokenizer
model_name = 'confidence_model'
model_path = model_name
tokenizer_path = model_name
tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_path)

model = ConfidenceEnhancedTransformer.from_pretrained(model_path, attn_implementation="eager")

# Move model to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Set the model to evaluation mode
model.eval()


main.py

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Model, GPT2Tokenizer
import torch.nn.functional as F

class ConfidenceEnhancedTransformer(GPT2LMHeadModel):
    def __init__(self, config):
        super(ConfidenceEnhancedTransformer, self).__init__(config)
        self.transformer = GPT2Model(config)
        #self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)  # Language modeling head

        # Confidence scoring head for epistemic uncertainty and OOD detection
        self.confidence_head = nn.Sequential(
            nn.Linear(config.n_embd, 128),
            nn.ReLU(),
            nn.Linear(128, 1),  # Single output for confidence score
            nn.Sigmoid()  # Confidence score between 0 and 1
        )

        # OOD detector head
        self.ood_detector = nn.Linear(config.n_embd, 1)  # Auxiliary head for OOD detection


train.py

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer, get_linear_schedule_with_warmup
from tqdm import tqdm
import torch.optim as optim
from datasets import load_dataset
from main import ConfidenceEnhancedTransformer
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import matplotlib.pyplot as plt

num_epochs = 3
save_steps = 500  # Save the model every 500 steps

# Define a custom dataset
class TextDataset(Dataset):
    def __init__(self, tokenizer, texts, block_size=128):
        self.examples = []

        for text in texts:


test.py

import unittest
import torch
from transformers import GPT2Tokenizer
from main import ConfidenceEnhancedTransformer

class TestConfidenceEnhancedTransformer(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        cls.model = ConfidenceEnhancedTransformer.from_pretrained('gpt2')
        cls.model.eval()  # Set model to evaluation mode

    def test_initialization(self):
        self.assertIsInstance(self.model, ConfidenceEnhancedTransformer)

    def test_forward_pass(self):
        input_text = "What is the capital of USA?"
        input_tokens = self.tokenizer(input_text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(input_tokens['input_ids'], num_dropout_samples=2)