Superposition Prompting: Improving and Accelerating Retrieval-Augmented Generation

pip install openai numpy

import os
import openai
import numpy as np
import asyncio
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
from typing import List, Dict, Tuple
from dataclasses import dataclass

# Set up OpenAI client
client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Data structures
@dataclass
class Document:
    title: str
    content: str

@dataclass
class PromptPath:
    preamble: str
    document: Document
    query: str

# Simulated KV cache
class KVCache:
    def __init__(self):
        self.cache = {}

    def set(self, key, value):
        self.cache[key] = value

    def get(self, key):
        return self.cache.get(key)

kv_cache = KVCache()

# Utility functions
def compute_logits(text: str) -> np.ndarray:
    """
    Simulate logit computation using OpenAI API.
    In a real implementation, this would be done by the language model itself.
    """
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": text}],
        max_tokens=1,
        logprobs=True,
        top_logprobs=5
    )
    # Extract logprobs from the response
    logprobs = response.choices[0].logprobs.content[0].top_logprobs
    return np.array([logprob.logprob for logprob in logprobs])

@lru_cache(maxsize=None)
def cached_compute_logits(text: str) -> np.ndarray:
    """
    Cached version of logit computation, implementing path caching technique.
    This reduces redundant computations for repeated text segments.
    """
    cached_result = kv_cache.get(text)
    if cached_result is not None:
        return cached_result
    result = compute_logits(text)
    kv_cache.set(text, result)
    return result

async def parallel_compute_bayesian_score(path: PromptPath) -> float:
    """
    Compute Bayesian score for a prompt path in parallel.
    This implements the path parallelization technique from the paper,
    allowing for efficient processing of multiple paths simultaneously.
    """
    loop = asyncio.get_event_loop()
    with ThreadPoolExecutor() as pool:
        preamble_logits = await loop.run_in_executor(pool, cached_compute_logits, path.preamble)
        document_logits = await loop.run_in_executor(pool, cached_compute_logits, path.document.content)
        query_logits = await loop.run_in_executor(pool, cached_compute_logits, path.query)

    # Compute Bayesian score components
    p_d_given_q = np.mean(document_logits)
    p_q_given_d = np.mean(query_logits)
    p_d = np.mean(preamble_logits)
    
    score = p_d_given_q + p_q_given_d + p_d
    print(f"Bayesian score for document '{path.document.title}': {score}")
    return score

def equilibrium_position(sequences: List[str]) -> List[List[float]]:
    """
    Implement the equilibrium position assignment technique.
    This assigns token positions to maintain consistent spacing across different length sequences.
    """
    total_length = sum(len(seq) for seq in sequences)
    average_length = total_length / len(sequences)
    positions = []
    current_position = 0
    for seq in sequences:
        seq_positions = [current_position + i * (average_length / len(seq)) for i in range(len(seq))]
        positions.append(seq_positions)
        current_position += average_length
    return positions

async def iterative_superposition(preamble: str, documents: List[Document], query: str, k: int, t: int) -> str:
    """
    Implement the iterative superposition technique for multi-hop reasoning.
    This function performs t iterations of path selection and query refinement.
    """
    print(f"\nStarting iterative superposition with k={k} and t={t}")
    paths = [PromptPath(preamble, doc, query) for doc in documents]
    for iteration in range(t):
        print(f"\nIteration {iteration + 1}:")
        # HIGHLIGHT: Parallel computation of Bayesian scores
        scores = await asyncio.gather(*[parallel_compute_bayesian_score(path) for path in paths])
        
        # HIGHLIGHT: Path pruning
        top_k_indices = np.argsort(scores)[-k:]
        paths = [paths[i] for i in top_k_indices]
        print(f"Selected top {k} documents: {[path.document.title for path in paths]}")
        
        # HIGHLIGHT: Multi-hop reasoning
        combined_content = " ".join([path.document.content for path in paths])
        query = f"Based on the following information: {combined_content}, {query}"
        paths = [PromptPath(preamble, path.document, query) for path in paths]
    
    # HIGHLIGHT: Equilibrium position assignment
    sequences = [preamble] + [path.document.content for path in paths] + [query]
    positions = equilibrium_position(sequences)
    
    final_prompt = ""
    for seq, pos in zip(sequences, positions):
        final_prompt += f"<tokens positions={pos}>{seq}</tokens>\n"
    
    print("\nGenerating final response...")
    # Generate response using OpenAI API
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a helpful assistant that answers questions based on the provided context."},
            {"role": "user", "content": final_prompt}
        ],
        max_tokens=150
    )
    
    return response.choices[0].message.content.strip()

async def superposition_prompting(preamble: str, documents: List[Document], query: str, k: int, t: int) -> str:
    """
    Main superposition prompting function that combines all techniques.
    """
    return await iterative_superposition(preamble, documents, query, k, t)

async def run_examples():
    """
    Run example queries to demonstrate the superposition prompting technique.
    """
    preamble = "You are a helpful assistant. Answer the following question based on the provided documents:"
    
    # HIGHLIGHT: More extensive RAG examples
    documents = [
        Document("Climate Change Overview", "Climate change refers to long-term shifts in temperatures and weather patterns. These shifts may be natural, but since the 1800s, human activities have been the main driver of climate change, primarily due to the burning of fossil fuels (like coal, oil, and gas), which produces heat-trapping gases."),
        Document("Greenhouse Effect", "The greenhouse effect is the way in which heat is trapped close to Earth's surface by 'greenhouse gases.' These heat-trapping gases can be thought of as a blanket wrapped around Earth, keeping the planet warmer than it would be without them."),
        Document("Carbon Dioxide Emissions", "Carbon dioxide (CO2) is the primary greenhouse gas emitted through human activities. In 2019, CO2 accounted for about 80 percent of all U.S. greenhouse gas emissions from human activities. The main human activity that emits CO2 is the combustion of fossil fuels for energy and transportation."),
        Document("Global Temperature Rise", "The planet's average surface temperature has risen about 2.12 degrees Fahrenheit (1.18 degrees Celsius) since the late 19th century, a change driven largely by increased carbon dioxide emissions into the atmosphere and other human activities."),
        Document("Sea Level Rise", "Global sea level has risen about 8 inches since reliable record keeping began in 1880. It is projected to rise another 1 to 8 feet by 2100. This is the result of added water from melting land ice and the expansion of seawater as it warms."),
        Document("Ocean Acidification", "Since the beginning of the Industrial Revolution, the acidity of surface ocean waters has increased by about 30 percent. This increase is the result of humans emitting more carbon dioxide into the atmosphere and hence more being absorbed into the ocean."),
        Document("Extreme Weather Events", "Climate change is causing more frequent and severe weather events, such as heat waves, droughts, and hurricanes. The number of record high temperature events in the United States has been increasing, while the number of record low temperature events has been decreasing, since 1950."),
        Document("Arctic Sea Ice Decline", "Both the extent and thickness of Arctic sea ice has declined rapidly over the last several decades. Arctic sea ice reaches its minimum each September. September Arctic sea ice is now declining at a rate of 13.1 percent per decade, relative to the 1981 to 2010 average."),
        Document("Glacier Retreat", "Glaciers are retreating almost everywhere around the world — including in the Alps, Himalayas, Andes, Rockies, Alaska, and Africa. Glacier National Park in Montana, USA, has lost over 120 glaciers in the last century."),
        Document("Biodiversity Loss", "Climate change is accelerating biodiversity loss across the globe. As temperatures change, many species are forced to migrate to new areas or face extinction. This disrupts ecosystems and food chains, potentially leading to cascading effects throughout the natural world.")
    ]
    
    queries = [
        "What are the main causes and effects of climate change?",
        "How does the greenhouse effect contribute to global warming?",
        "What are the projected consequences of sea level rise and ocean acidification?",
        "How is climate change affecting weather patterns and biodiversity?",
        "What evidence supports the claim that human activities are the main driver of recent climate change?"
    ]

    for query in queries:
        print(f"\n\nQuery: {query}")
        answer = await superposition_prompting(preamble, documents, query, k=3, t=2)
        print(f"Answer: {answer}")
        print("-" * 80)

# Run the examples
if __name__ == "__main__":
    asyncio.run(run_examples())