Question Answering Transformers: Last Chars Bug

import torch
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
from torch.nn.functional import softmax

# Load a pre-fine-tuned QA model (BERT-base on SQuAD v1.1)
# For production, fine-tune on your own domain data first.
model_name = "distilbert-base-cased-distilled-squad"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

# Input examples
context = """
The Transformer architecture was introduced in the 2017 paper "Attention Is All You Need"
by Vaswani et al. from Google Brain and the University of Toronto. It has since become
foundational for most state-of-the-art NLP models including BERT, GPT, and T5.
"""

question = "What paper introduced the Transformer architecture?"

# Tokenise with offsets for character-level alignment
inputs = tokenizer(
    question,
    context,
    return_tensors="pt",           # return PyTorch tensors
    return_offsets_mapping=True,   # IMPORTANT: get character positions for each token
    truncation=True,
    max_length=512
)

offset_mapping = inputs.pop("offset_mapping").squeeze().tolist()

with torch.no_grad():
    outputs = model(**inputs)

start_logits = outputs.start_logits.squeeze()
end_logits = outputs.end_logits.squeeze()

# Find best start and end positions
start_idx = torch.argmax(start_logits).item()
end_idx = torch.argmax(end_logits).item()

# Extract answer using offset mapping (corrects for subword tokens!)
if start_idx <= end_idx and start_idx < len(offset_mapping) and end_idx < len(offset_mapping):
    start_char = offset_mapping[start_idx][0]
    end_char = offset_mapping[end_idx][1]
    answer = context[start_char:end_char]
    confidence = softmax(start_logits)[start_idx].item() * softmax(end_logits)[end_idx].item()
    print(f"Question: {question}")
    print(f"Answer: {answer}")
    print(f"Confidence: {confidence:.4f}")
    print(f"Span: chars {start_char}-{end_char}")
else:
    print("No answer found (invalid span)")

# Example of subword tokenisation issue (if we had concatenated tokens instead of using offsets)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze())
print(f"\nToken sequence: {tokens[start_idx:end_idx+1]}")
print(f"Naive concatenation would give: '{''.join(tokens[start_idx:end_idx+1]).replace('##', '')}'")
print(f"Offset mapping gives correct: '{answer}'")

# Confidence thresholding for SQuAD2.0 style (no answer)
null_score = start_logits[0] + end_logits[0]  # [CLS] token position
best_span_score = start_logits[start_idx] + end_logits[end_idx]

print(f"\nNull answer score: {null_score:.4f}")
print(f"Best span score: {best_span_score:.4f}")

if best_span_score < null_score + 0.5:  # threshold tuned on dev set
    print("Model would predict NO ANSWER (thresholded)")

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForQuestionAnswering, TrainingArguments, Trainer
from transformers import DefaultDataCollator
import numpy as np

# Load SQuAD2.0 dataset
# For production, replace with your own dataset in the same format
squad = load_dataset("squad_v2")
train_dataset = squad["train"]
valid_dataset = squad["validation"]

model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForQuestionAnswering.from_pretrained(model_name)

# Preprocessing function for QA
def preprocess_qa(examples):
    questions = [q.strip() for q in examples["question"]]
    contexts = examples["context"]
    
    # Tokenise with offset mapping to find answer span positions
    tokenized = tokenizer(
        questions,
        contexts,
        truncation="only_second",  # only truncate context, preserve question
        max_length=384,
        stride=128,               # overlap for contexts that exceed max_length
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )
    
    sample_mapping = tokenized.pop("overflow_to_sample_mapping")
    offset_mapping = tokenized.pop("offset_mapping")
    
    start_positions = []
    end_positions = []
    
    for i, offsets in enumerate(offset_mapping):
        sample_idx = sample_mapping[i]
        answer = examples["answer"][sample_idx]
        is_impossible = examples["is_impossible"][sample_idx]
        
        if not is_impossible:
            answer_start_char = answer["answer_start"][0]
            answer_text = answer["text"][0]
            end_char = answer_start_char + len(answer_text)
            
            # Find start token index
            start_idx = None
            end_idx = None
            for idx, (start_char, end_char_token) in enumerate(offsets):
                if start_char <= answer_start_char < end_char_token:
                    start_idx = idx
                if start_char < end_char <= end_char_token:
                    end_idx = idx
            
            if start_idx is not None and end_idx is not None:
                start_positions.append(start_idx)
                end_positions.append(end_idx)
            else:
                start_positions.append(0)
                end_positions.append(0)
        else:
            # Unanswerable questions: answer is CLS token (index 0)
            start_positions.append(0)
            end_positions.append(0)
    
    tokenized["start_positions"] = start_positions
    tokenized["end_positions"] = end_positions
    return tokenized

# Apply preprocessing
train_tokenized = train_dataset.map(preprocess_qa, batched=True, remove_columns=train_dataset.column_names)
valid_tokenized = valid_dataset.map(preprocess_qa, batched=True, remove_columns=valid_dataset.column_names)

# Training arguments
training_args = TrainingArguments(
    output_dir="./qa-model",
    evaluation_strategy="epoch",
    learning_rate=3e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir="./logs",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="eval_f1",
)

data_collator = DefaultDataCollator()

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_tokenized,
    eval_dataset=valid_tokenized,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

# Train the model
# trainer.train()

# Save for production
# model.save_pretrained("./production-qa-model")
# tokenizer.save_pretrained("./production-qa-model")

print("Training pipeline configured. Uncomment trainer.train() to run.")

import torch
from transformers import AutoTokenizer, AutoModelForQuestionAnswering

def qa_with_sliding_window(
    question: str,
    context: str,
    model,
    tokenizer,
    max_length: int = 384,
    stride: int = 128,
    threshold: float = 0.0
):
    """
    Run QA inference on long contexts using a sliding window.
    Returns the highest-scoring answer span across all windows.
    """
    # Tokenise without truncation to see how many tokens we have
    full_tokens = tokenizer(question, context, truncation=False, return_tensors="pt")
    total_tokens = full_tokens["input_ids"].shape[1]
    
    if total_tokens <= max_length:
        # Short context: normal inference
        inputs = tokenizer(question, context, return_tensors="pt", return_offsets_mapping=True, max_length=max_length, truncation=True)
        with torch.no_grad():
            outputs = model(**{k: v for k, v in inputs.items() if k != "offset_mapping"})
        start_logits = outputs.start_logits.squeeze()
        end_logits = outputs.end_logits.squeeze()
        start_idx = torch.argmax(start_logits).item()
        end_idx = torch.argmax(end_logits).item()
        score = start_logits[start_idx].item() + end_logits[end_idx].item()
        offset_mapping = inputs["offset_mapping"].squeeze().tolist()
        if start_idx <= end_idx and score > threshold:
            return extract_answer_from_offsets(context, offset_mapping, start_idx, end_idx), score
        return None, score
    
    # Long context: sliding window
    best_answer = None
    best_score = -float("inf")
    
    # Tokenise with stride
    inputs = tokenizer(
        question, context,
        return_tensors="pt",
        return_offsets_mapping=True,
        max_length=max_length,
        stride=stride,
        truncation="only_second",
        return_overflowing_tokens=True
    )
    
    offset_mappings = inputs.pop("offset_mapping")
    for i, offset_mapping in enumerate(offset_mappings):
        # Prepare single-window inputs
        window_inputs = {k: v[i].unsqueeze(0) for k, v in inputs.items() if k != "offset_mapping"}
        
        with torch.no_grad():
            outputs = model(**window_inputs)
        
        start_logits = outputs.start_logits.squeeze()
        end_logits = outputs.end_logits.squeeze()
        start_idx = torch.argmax(start_logits).item()
        end_idx = torch.argmax(end_logits).item()
        score = start_logits[start_idx].item() + end_logits[end_idx].item()
        
        # Skip if span goes beyond this window's valid tokens
        if start_idx <= end_idx and start_idx < len(offset_mapping) and end_idx < len(offset_mapping):
            # Check if this window's answer is better than previous
            if score > best_score:
                best_score = score
                best_answer = extract_answer_from_offsets(context, offset_mapping.tolist(), start_idx, end_idx)
    
    return best_answer, best_score

def extract_answer_from_offsets(context, offset_mapping, start_idx, end_idx):
    """Extract answer using offset mapping (fixes subword issues)."""
    start_char = offset_mapping[start_idx][0]
    end_char = offset_mapping[end_idx][1]
    return context[start_char:end_char]

# Usage
# model = AutoModelForQuestionAnswering.from_pretrained("bert-base-uncased")
# tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
# long_context = "..." * 10000
# answer, score = qa_with_sliding_window("What is the main finding?", long_context, model, tokenizer)

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForQuestionAnswering
from optimize import ORTModelForQuestionAnswering  # requires optimum

def quantize_to_int8(model, model_name: str, save_path: str):
    """Convert FP32 model to INT8 quantized for CPU inference."""
    from optimum.onnxruntime import ORTModelForQuestionAnswering
    from optimum.onnxruntime.configuration import AutoQuantizationConfig
    
    # Export to ONNX
    ort_model = ORTModelForQuestionAnswering.from_pretrained(
        model_name, export=True, provider="CPUExecutionProvider"
    )
    
    # Apply dynamic quantization
    qconfig = AutoQuantizationConfig.avx512_vnni(is_static=False)
    quantized_model = ORTModelForQuestionAnswering.from_pretrained(
        model_name,
        quantization_config=qconfig,
        export=True,
        provider="CPUExecutionProvider"
    )
    quantized_model.save_pretrained(save_path)
    print(f"INT8 quantised model saved to {save_path}")
    return quantized_model

def tune_null_threshold(model, tokenizer, validation_dataset):
    """Find optimal null threshold for unanswerable questions."""
    score_diffs = []  # max_span_score - null_score
    has_answer = []   # ground truth
    
    for ex in validation_dataset:
        inputs = tokenizer(ex["question"], ex["context"], return_tensors="pt", truncation=True, max_length=384)
        with torch.no_grad():
            outputs = model(**inputs)
        
        start_logits = outputs.start_logits.squeeze()
        end_logits = outputs.end_logits.squeeze()
        
        # Best span excluding CLS
        start_logits_no_cls = start_logits[1:]
        end_logits_no_cls = end_logits[1:]
        best_span_score = (start_logits_no_cls.unsqueeze(1) + end_logits_no_cls.unsqueeze(0)).max().item()
        
        null_score = start_logits[0].item() + end_logits[0].item()
        score_diffs.append(best_span_score - null_score)
        has_answer.append(not ex["is_impossible"])
    
    score_diffs = np.array(score_diffs)
    has_answer = np.array(has_answer)
    
    # Find thresholds
    thresholds = np.percentile(score_diffs, np.linspace(0, 100, 101))
    best_f1 = 0
    best_threshold = 0
    
    for t in thresholds:
        predicted_has_answer = score_diffs > t
        tp = np.sum(predicted_has_answer & has_answer)
        fp = np.sum(predicted_has_answer & ~has_answer)
        fn = np.sum(~predicted_has_answer & has_answer)
        
        precision = tp / (tp + fp) if tp + fp > 0 else 0
        recall = tp / (tp + fn) if tp + fn > 0 else 0
        f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0
        
        if f1 > best_f1:
            best_f1 = f1
            best_threshold = t
    
    print(f"Optimal threshold: {best_threshold:.4f} (F1: {best_f1:.4f})")
    return best_threshold

def benchmark_latency(model, tokenizer, text_sample, num_runs=100):
    """Measure average inference latency."""
    import time
    inputs = tokenizer("What is the answer?", text_sample, return_tensors="pt", truncation=True, max_length=384)
    
    # Warmup
    for _ in range(10):
        with torch.no_grad():
            _ = model(**inputs)
    
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(num_runs):
        with torch.no_grad():
            _ = model(**inputs)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - start
    
    print(f"Latency: {elapsed/num_runs*1000:.2f} ms per query ({num_runs} runs)")
    return elapsed / num_runs

# Usage
# quantized = quantize_to_int8("bert-base-uncased", "./quantized-qa")
# threshold = tune_null_threshold(quantized, tokenizer, validation_dataset)
# latency_per_query = benchmark_latency(quantized, tokenizer, long_context)