Self-RAG in Machine Learning

"""Self-RAG inference pipeline with adaptive retrieval and reflection tokens."""
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from typing import List, Dict, Tuple, Optional
import numpy as np


class SelfRAGInference:
    """Inference engine for a Self-RAG fine-tuned model."""

    # Reflection token definitions
    RETRIEVE_TOKENS = {"[Retrieve=Yes]": True, "[Retrieve=No]": False, "[Retrieve=Continue]": None}
    ISREL_TOKENS = {"[IsRel=Relevant]": 1.0, "[IsRel=Irrelevant]": 0.0}
    ISSUP_TOKENS = {
        "[IsSup=Fully Supported]": 1.0,
        "[IsSup=Partially Supported]": 0.5,
        "[IsSup=No Support]": 0.0,
    }
    ISUSE_TOKENS = {f"[IsUse={i}]": i / 5.0 for i in range(1, 6)}

    def __init__(
        self,
        model_name: str = "selfrag/selfrag_llama2_7b",
        retriever=None,
        alpha: float = 1.0,
        beta: float = 1.0,
        gamma: float = 0.5,
        top_k: int = 5,
        max_segments: int = 10,
    ):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name, torch_dtype=torch.float16, device_map="auto"
        )
        self.retriever = retriever
        self.alpha = alpha   # weight for relevance
        self.beta = beta     # weight for support (factuality)
        self.gamma = gamma   # weight for utility
        self.top_k = top_k
        self.max_segments = max_segments

    def generate_with_reflection(self, query: str) -> Dict:
        """Generate a response with adaptive retrieval and self-reflection."""
        segments = []
        full_context = f"### Instruction:\n{query}\n### Response:\n"
        retrieval_count = 0

        for seg_idx in range(self.max_segments):
            # Step 1: Generate retrieve decision
            retrieve_decision = self._predict_retrieve_token(full_context)

            if retrieve_decision is True and self.retriever is not None:
                # Step 2: Retrieve passages
                retrieval_count += 1
                passages = self.retriever.search(query, top_k=self.top_k)

                # Step 3: Generate candidate segments for each passage
                candidates = []
                for passage in passages:
                    candidate = self._generate_candidate_segment(
                        full_context, passage
                    )
                    candidates.append(candidate)

                # Step 4: Score and select best candidate
                best = self._select_best_candidate(candidates)
                segments.append(best)
                full_context += best["text"]
            else:
                # Generate without retrieval
                segment_text = self._generate_segment(full_context)
                segments.append({"text": segment_text, "retrieved": False})
                full_context += segment_text

            # Check for end of generation
            if self._is_generation_complete(full_context):
                break

        return {
            "response": self._assemble_response(segments),
            "segments": segments,
            "retrieval_count": retrieval_count,
            "total_segments": len(segments),
        }

    def _predict_retrieve_token(self, context: str) -> Optional[bool]:
        """Predict whether retrieval is needed at this point."""
        inputs = self.tokenizer(context, return_tensors="pt").to(self.model.device)
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits[0, -1, :]

        # Get probabilities for retrieve tokens
        retrieve_probs = {}
        for token_str, value in self.RETRIEVE_TOKENS.items():
            token_id = self.tokenizer.encode(token_str, add_special_tokens=False)
            if token_id:
                retrieve_probs[value] = logits[token_id[0]].item()

        # Return the decision with highest logit
        return max(retrieve_probs, key=retrieve_probs.get)

    def _generate_candidate_segment(
        self, context: str, passage: Dict
    ) -> Dict:
        """Generate a candidate segment conditioned on a retrieved passage."""
        augmented_context = (
            f"{context}[Retrieve=Yes] "
            f"[Document] {passage['text']} [/Document]\n"
        )
        inputs = self.tokenizer(
            augmented_context, return_tensors="pt", truncation=True, max_length=2048
        ).to(self.model.device)

        with torch.no_grad():
            output_ids = self.model.generate(
                **inputs, max_new_tokens=256, do_sample=False
            )

        generated = self.tokenizer.decode(
            output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=False
        )

        # Parse reflection tokens from generated text
        isrel = self._extract_reflection_score(generated, self.ISREL_TOKENS)
        issup = self._extract_reflection_score(generated, self.ISSUP_TOKENS)
        isuse = self._extract_reflection_score(generated, self.ISUSE_TOKENS)

        # Clean text (remove reflection tokens)
        clean_text = self._strip_reflection_tokens(generated)

        return {
            "text": clean_text,
            "retrieved": True,
            "passage": passage["text"][:200],
            "scores": {"isrel": isrel, "issup": issup, "isuse": isuse},
        }

    def _select_best_candidate(self, candidates: List[Dict]) -> Dict:
        """Score candidates using weighted reflection tokens."""
        best_score = -float("inf")
        best_candidate = candidates[0]

        for candidate in candidates:
            s = candidate["scores"]
            score = (
                self.alpha * s["isrel"]
                + self.beta * s["issup"]
                + self.gamma * s["isuse"]
            )
            if score > best_score:
                best_score = score
                best_candidate = candidate

        best_candidate["beam_score"] = best_score
        return best_candidate

    def _extract_reflection_score(self, text: str, token_map: Dict) -> float:
        """Extract the reflection token score from generated text."""
        for token_str, score in token_map.items():
            if token_str in text:
                return score
        return 0.0  # default if no reflection token found

    def _strip_reflection_tokens(self, text: str) -> str:
        """Remove all reflection tokens from text."""
        import re
        pattern = r'\[(?:Retrieve|IsRel|IsSup|IsUse)=[^\]]*\]'
        return re.sub(pattern, '', text).strip()

    def _generate_segment(self, context: str) -> str:
        """Generate a text segment without retrieval."""
        inputs = self.tokenizer(
            context, return_tensors="pt", truncation=True, max_length=2048
        ).to(self.model.device)
        with torch.no_grad():
            output_ids = self.model.generate(
                **inputs, max_new_tokens=256, do_sample=False
            )
        text = self.tokenizer.decode(
            output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=False
        )
        return self._strip_reflection_tokens(text)

    def _is_generation_complete(self, context: str) -> bool:
        """Check if generation should stop."""
        return context.rstrip().endswith(self.tokenizer.eos_token) or len(context) > 8000

    def _assemble_response(self, segments: List[Dict]) -> str:
        """Combine segments into final response."""
        return " ".join(seg["text"] for seg in segments).strip()

"""Generate reflection token labels for Self-RAG training data using a critic model."""
import json
from openai import AzureOpenAI
from typing import List, Dict
from dataclasses import dataclass, asdict
from enum import Enum


class RetrieveLabel(Enum):
    YES = "Yes"
    NO = "No"
    CONTINUE = "Continue"


class IsRelLabel(Enum):
    RELEVANT = "Relevant"
    IRRELEVANT = "Irrelevant"


class IsSupLabel(Enum):
    FULLY_SUPPORTED = "Fully Supported"
    PARTIALLY_SUPPORTED = "Partially Supported"
    NO_SUPPORT = "No Support"


@dataclass
class ReflectionLabels:
    retrieve: str
    isrel: str = None
    issup: str = None
    isuse: int = None


class SelfRAGCritic:
    """Critic model that labels training data with reflection tokens."""

    def __init__(self, client: AzureOpenAI, deployment: str = "gpt-4"):
        self.client = client
        self.deployment = deployment

    def label_retrieve(self, query: str, partial_response: str) -> str:
        """Determine if retrieval is needed at this generation point."""
        prompt = f"""Given the following query and partial response, determine
whether retrieving external information would improve the response.

Query: {query}
Partial Response: {partial_response}

Classify as one of:
- "Yes": External knowledge is needed for a factual, complete answer
- "No": The model can answer confidently from parametric knowledge
- "Continue": The current segment does not require new retrieval

Output ONLY the classification label."""

        response = self.client.chat.completions.create(
            model=self.deployment,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=10,
        )
        label = response.choices[0].message.content.strip().strip('"')
        return label if label in [e.value for e in RetrieveLabel] else "No"

    def label_relevance(self, query: str, passage: str) -> str:
        """Judge whether a retrieved passage is relevant to the query."""
        prompt = f"""Given the query and retrieved passage, determine relevance.

Query: {query}
Passage: {passage[:1000]}

Is this passage relevant to answering the query?
Output ONLY: "Relevant" or "Irrelevant""""

        response = self.client.chat.completions.create(
            model=self.deployment,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=10,
        )
        label = response.choices[0].message.content.strip().strip('"')
        return label if label in [e.value for e in IsRelLabel] else "Irrelevant"

    def label_support(self, response_segment: str, passage: str) -> str:
        """Verify if the response segment is supported by the passage."""
        prompt = f"""Given the response segment and the passage it was based on,
determine the level of factual support.

Response Segment: {response_segment}
Passage: {passage[:1000]}

Classify support level as one of:
- "Fully Supported": All claims in the response are backed by the passage
- "Partially Supported": Some claims are supported, others are not
- "No Support": The response is not grounded in the passage

Output ONLY the classification label."""

        response = self.client.chat.completions.create(
            model=self.deployment,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=20,
        )
        label = response.choices[0].message.content.strip().strip('"')
        return label if label in [e.value for e in IsSupLabel] else "No Support"

    def label_utility(self, query: str, full_response: str) -> int:
        """Rate overall response utility on a 1-5 scale."""
        prompt = f"""Rate the utility of this response for the given query.

Query: {query}
Response: {full_response[:2000]}

Rate from 1 to 5:
1 = Completely unhelpful or wrong
2 = Mostly unhelpful with major gaps
3 = Partially helpful but incomplete
4 = Helpful with minor issues
5 = Excellent, comprehensive, and accurate

Output ONLY the numeric rating."""

        response = self.client.chat.completions.create(
            model=self.deployment,
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=5,
        )
        try:
            score = int(response.choices[0].message.content.strip())
            return max(1, min(5, score))
        except ValueError:
            return 3  # default to neutral

    def label_training_example(
        self, query: str, response: str, passages: List[str]
    ) -> Dict:
        """Label a complete training example with all reflection tokens."""
        retrieve_label = self.label_retrieve(query, "")
        labels = {"query": query, "response": response, "retrieve": retrieve_label}

        if retrieve_label == "Yes" and passages:
            passage_labels = []
            for passage in passages:
                isrel = self.label_relevance(query, passage)
                issup = self.label_support(response, passage)
                passage_labels.append(
                    {"passage": passage[:500], "isrel": isrel, "issup": issup}
                )
            labels["passage_labels"] = passage_labels

        labels["isuse"] = self.label_utility(query, response)
        return labels


# Usage
def prepare_training_data(raw_data: List[Dict], output_path: str):
    """Label a dataset with reflection tokens for Self-RAG training."""
    client = AzureOpenAI(
        azure_endpoint="https://your-resource.openai.azure.com",
        api_key="your-key",
        api_version="2024-02-15-preview",
    )
    critic = SelfRAGCritic(client)
    labeled_data = []

    for item in raw_data:
        labeled = critic.label_training_example(
            query=item["query"],
            response=item["response"],
            passages=item.get("passages", []),
        )
        labeled_data.append(labeled)

    with open(output_path, "w") as f:
        for item in labeled_data:
            f.write(json.dumps(item) + "\n")

    print(f"Labeled {len(labeled_data)} examples -> {output_path}")

"""Fine-tune an LLM to generate reflection tokens (Self-RAG generator training)."""
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
)
from datasets import Dataset
import json
from typing import List, Dict


SPECIAL_TOKENS = [
    "[Retrieve=Yes]", "[Retrieve=No]", "[Retrieve=Continue]",
    "[IsRel=Relevant]", "[IsRel=Irrelevant]",
    "[IsSup=Fully Supported]", "[IsSup=Partially Supported]", "[IsSup=No Support]",
    "[IsUse=1]", "[IsUse=2]", "[IsUse=3]", "[IsUse=4]", "[IsUse=5]",
    "[Document]", "[/Document]",
]


def prepare_tokenizer(model_name: str) -> AutoTokenizer:
    """Add reflection tokens to the tokenizer vocabulary."""
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.add_special_tokens({"additional_special_tokens": SPECIAL_TOKENS})
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    return tokenizer


def format_training_example(item: Dict) -> str:
    """Convert a critic-labeled example into a training sequence."""
    parts = [f"### Instruction:\n{item['query']}\n### Response:\n"]

    if item["retrieve"] == "Yes" and "passage_labels" in item:
        parts.append(f"[Retrieve=Yes]")
        # Use the best passage (highest support)
        best_passage = max(
            item["passage_labels"],
            key=lambda p: 1.0 if p["issup"] == "Fully Supported" else 0.0,
        )
        parts.append(f"[IsRel={best_passage['isrel']}]")
        parts.append(f"[Document] {best_passage['passage']} [/Document]")
        parts.append(item["response"])
        parts.append(f"[IsSup={best_passage['issup']}]")
    else:
        parts.append(f"[Retrieve=No]")
        parts.append(item["response"])

    parts.append(f"[IsUse={item['isuse']}]")
    return " ".join(parts)


def train_self_rag_generator(
    model_name: str = "meta-llama/Llama-2-7b-hf",
    training_data_path: str = "labeled_data.jsonl",
    output_dir: str = "./selfrag-model",
    epochs: int = 3,
    batch_size: int = 4,
    learning_rate: float = 2e-5,
    max_length: int = 2048,
):
    """Fine-tune an LLM to become a Self-RAG generator."""
    # 1. Load and prepare tokenizer
    tokenizer = prepare_tokenizer(model_name)

    # 2. Load model and resize embeddings for new tokens
    model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, device_map="auto"
    )
    model.resize_token_embeddings(len(tokenizer))

    # 3. Load and format training data
    with open(training_data_path) as f:
        raw_data = [json.loads(line) for line in f]

    formatted_texts = [format_training_example(item) for item in raw_data]

    # 4. Tokenize
    def tokenize_fn(examples):
        encodings = tokenizer(
            examples["text"],
            truncation=True,
            max_length=max_length,
            padding="max_length",
        )
        encodings["labels"] = encodings["input_ids"].copy()
        return encodings

    dataset = Dataset.from_dict({"text": formatted_texts})
    tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=["text"])

    # 5. Training arguments
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=4,
        learning_rate=learning_rate,
        warmup_ratio=0.05,
        weight_decay=0.01,
        logging_steps=50,
        save_strategy="epoch",
        bf16=True,
        gradient_checkpointing=True,
        report_to="wandb",
    )

    # 6. Train
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized,
        tokenizer=tokenizer,
    )
    trainer.train()

    # 7. Save
    trainer.save_model(output_dir)
    tokenizer.save_pretrained(output_dir)
    print(f"Self-RAG generator saved to {output_dir}")

"""Controllable Self-RAG inference — adjust factuality vs creativity at runtime."""
from dataclasses import dataclass
from typing import List, Dict, Optional
import numpy as np


@dataclass
class InferenceConfig:
    """Runtime configuration for Self-RAG inference."""
    alpha: float = 1.0    # relevance weight
    beta: float = 1.0     # support/factuality weight
    gamma: float = 0.5    # utility weight
    top_k: int = 5        # passages to retrieve
    retrieval_threshold: float = 0.5  # min probability for Retrieve=Yes
    min_support_threshold: float = 0.3  # min IsSup score to accept a segment
    max_segments: int = 10


# Preset configurations for different use cases
PRESETS = {
    "factual_qa": InferenceConfig(
        alpha=1.0, beta=2.0, gamma=0.3,
        retrieval_threshold=0.3,  # retrieve more aggressively
        min_support_threshold=0.7,  # require high support
    ),
    "creative_writing": InferenceConfig(
        alpha=0.3, beta=0.3, gamma=2.0,
        retrieval_threshold=0.8,  # retrieve less often
        min_support_threshold=0.0,  # allow unsupported content
    ),
    "balanced": InferenceConfig(
        alpha=1.0, beta=1.0, gamma=1.0,
        retrieval_threshold=0.5,
        min_support_threshold=0.3,
    ),
    "citation_heavy": InferenceConfig(
        alpha=1.5, beta=2.5, gamma=0.5,
        top_k=10,
        retrieval_threshold=0.2,  # almost always retrieve
        min_support_threshold=0.8,  # very strict support
    ),
}


class ControllableSelfRAG:
    """Self-RAG with runtime-adjustable behavior."""

    def __init__(self, model, retriever, config: InferenceConfig = None):
        self.model = model
        self.retriever = retriever
        self.config = config or InferenceConfig()

    def set_preset(self, preset_name: str):
        """Switch behavior preset at runtime — no retraining needed."""
        if preset_name not in PRESETS:
            raise ValueError(f"Unknown preset: {preset_name}. Available: {list(PRESETS.keys())}")
        self.config = PRESETS[preset_name]
        return self.config

    def score_candidate(self, scores: Dict[str, float]) -> float:
        """Score a candidate segment using current config weights."""
        return (
            self.config.alpha * scores.get("isrel", 0.0)
            + self.config.beta * scores.get("issup", 0.0)
            + self.config.gamma * scores.get("isuse", 0.0)
        )

    def should_retrieve(self, retrieve_prob: float) -> bool:
        """Decide whether to retrieve based on threshold."""
        return retrieve_prob >= self.config.retrieval_threshold

    def should_accept_segment(self, issup_score: float) -> bool:
        """Decide whether a segment meets minimum support requirements."""
        return issup_score >= self.config.min_support_threshold

    def generate(self, query: str) -> Dict:
        """Generate response with current configuration."""
        segments = []
        retrieval_events = []

        for seg_idx in range(self.config.max_segments):
            # Check retrieve probability
            retrieve_prob = self.model.get_retrieve_probability(query, segments)

            if self.should_retrieve(retrieve_prob):
                passages = self.retriever.search(query, top_k=self.config.top_k)
                candidates = []

                for passage in passages:
                    candidate = self.model.generate_segment(query, segments, passage)
                    candidate["beam_score"] = self.score_candidate(candidate["scores"])
                    candidates.append(candidate)

                # Sort by beam score and filter by support threshold
                candidates.sort(key=lambda c: c["beam_score"], reverse=True)
                accepted = [
                    c for c in candidates
                    if self.should_accept_segment(c["scores"].get("issup", 0))
                ]

                if accepted:
                    best = accepted[0]
                else:
                    # Fallback: take highest-scored even if below support threshold
                    best = candidates[0]
                    best["support_warning"] = True

                segments.append(best)
                retrieval_events.append({
                    "segment": seg_idx,
                    "passages_considered": len(passages),
                    "accepted_candidates": len(accepted),
                    "best_score": best["beam_score"],
                })
            else:
                segment = self.model.generate_segment(query, segments, passage=None)
                segments.append(segment)

            if self.model.is_complete(segments):
                break

        return {
            "response": " ".join(s.get("text", "") for s in segments),
            "config": {
                "alpha": self.config.alpha,
                "beta": self.config.beta,
                "gamma": self.config.gamma,
            },
            "retrieval_events": retrieval_events,
            "total_segments": len(segments),
            "warnings": [
                s for s in segments if s.get("support_warning")
            ],
        }


# Usage example
def demo_controllable_inference():
    """Show how the same model behaves differently with different configs."""
    # model and retriever would be initialized here
    rag = ControllableSelfRAG(model=None, retriever=None)

    # For medical QA — maximize factuality
    rag.set_preset("factual_qa")
    # result = rag.generate("What are the side effects of metformin?")

    # For story writing — maximize creativity
    rag.set_preset("creative_writing")
    # result = rag.generate("Write a short story about a robot learning to paint")

    # For research — maximize citations
    rag.set_preset("citation_heavy")
    # result = rag.generate("Summarize recent advances in protein folding")

# Self-RAG Inference Configuration (YAML)
model:
  name: selfrag/selfrag_llama2_13b
  dtype: float16
  device_map: auto
  max_length: 2048

retriever:
  type: contriever
  index_path: /data/indices/wiki_2023
  top_k: 5
  batch_size: 32

reflection_weights:
  alpha: 1.0          # IsRel weight (relevance)
  beta: 1.5           # IsSup weight (factuality)
  gamma: 0.5          # IsUse weight (utility)

inference:
  max_segments: 10
  segment_max_tokens: 256
  retrieval_threshold: 0.5    # min prob for Retrieve=Yes
  min_support_score: 0.3      # min IsSup to accept segment
  strip_reflection_tokens: true
  beam_width: 1               # 1 = greedy, >1 = beam search over segments

presets:
  factual_qa:
    alpha: 1.0
    beta: 2.0
    gamma: 0.3
    retrieval_threshold: 0.3
  creative:
    alpha: 0.3
    beta: 0.3
    gamma: 2.0
    retrieval_threshold: 0.8
  citation_heavy:
    alpha: 1.5
    beta: 2.5
    gamma: 0.5
    retrieval_threshold: 0.2