Project Overview

This project implements a hybrid Retrieval-Augmented Generation (RAG) system that combines structured database queries with unstructured document retrieval. The application allows users to ask natural language questions about US cities and receive accurate, contextual responses. The system uses a dual-pipeline approach:

1. Structured Data Pipeline: SQL database containing city statistics (population, state)

2. Unstructured Data Pipeline: LlamaCloud for document retrieval of general city information

The application is built with a Streamlit frontend that provides an intuitive chat interface, allowing users to interact with the system conversationally.

Architecture Details

Component Breakdown

1. User Interface Layer

• Technology: Streamlit

• Purpose: Provides chat interface, API key configuration, and displays responses

• Key Features:

  • Chat history tracking using session state

  • Sidebar configuration for API keys and LlamaCloud settings

  • Real-time response generation with loading indicators

2. Structured Data Layer

• Technology: SQLite (in-memory), SQLAlchemy

• Purpose: Stores and queries factual city information

• Components:

  • setup_database(): Initializes the database with city data

  • CityQueryEngine: Custom class for direct SQL querying

  • NLSQLTableQueryEngine: LlamaIndex component for natural language to SQL conversion

3. Unstructured Data Layer

• Technology: LlamaCloud

• Purpose: Stores and retrieves document-based city information

• Components:

  • LlamaCloudIndex: Connects to pre-indexed documents in LlamaCloud

  • Vector query engine for semantic retrieval

4. LLM Integration Layer

• Technology: Gemini 2.0 Flash

• Purpose: Natural language understanding and response generation

• Components:

  • Text generation model for synthesizing responses

  • Embedding model for vector search (when needed)

5. Query Orchestration Layer

• Purpose: Routes queries to appropriate backends based on content

• Logic:

  • Population/statistic queries → SQL pipeline

  • General information queries → LlamaCloud pipeline

Data Flow

1. Input Phase:

   • User submits query through Streamlit chat interface

   • Query is added to chat history

   • System analyzes query content

2. Processing Phase:

   • If query contains population keywords:

     • First attempt direct SQL via pattern matching

     • Fall back to NLSQLTableQueryEngine if needed

   • If query is about general information:

     • Route to LlamaCloud vector search

   • Response is generated using retrieved context

3. Output Phase:

   • Response displayed in chat interface

   • Source of information (Database or LlamaCloud) is displayed

   • Response added to chat history

Implementation Details

Database Schema

• The application uses a simple SQLite database with a single table:

CREATE TABLE city_stats (
    city_name VARCHAR(16) PRIMARY KEY,
    population INTEGER,
    state VARCHAR(16) NOT NULL
);

Data is populated with information about six major US cities:

• New York City (New York)

• Los Angeles (California)

• Chicago (Illinois)

• Houston (Texas)

• Miami (Florida)

• Seattle (Washington)

SQL Query Engine

The custom CityQueryEngine class provides specialized query capabilities:

1. Direct Execution: execute_query() method runs raw SQL and formats results

2. Specialized Queries:

   • query_highest_population(): Finds city with highest population

   • query_lowest_population(): Finds city with lowest population

   • query_all_cities_ranked(): Orders all cities by population

   • query_by_state(): Filters cities by state name

3. Natural Language Processing:

   • process_population_query(): Uses regular expressions and keyword matching to convert natural language to appropriate SQL queries

   • Pattern matching for terms like "highest", "lowest", "in [state]", etc.

LlamaCloud Integration

The application connects to LlamaCloud for document retrieval:

1. Configuration:

   • API key

   • Organization ID

   • Project name

   • Index name

2. Index Connection:

   • Creates LlamaCloudIndex object

   • Configures vector query engine

3. Query Execution:

   • Sends natural language query directly to LlamaCloud

   • Retrieves relevant document chunks

   • Uses context for response generation

Model Configuration

The application uses Gemini 2.0 Flash for both understanding and generation:

1. LLM Model:

    gemini_model = Gemini(
        model="models/gemini-2.0-flash",
        api_key=st.session_state.GOOGLE_API_KEY,
        temperature=0.2
    )
   

2. Embedding Model:

    gemini_embed_model = GeminiEmbedding(
        model_name="models/embedding-001",
        api_key=st.session_state.GOOGLE_API_KEY
    )
   

3. Global Settings:

    Settings.llm = gemini_model
    Settings.embed_model = gemini_embed_model
   

Query Processing Logic

The core logic that determines how to handle each query:

# Check if this is a population query
if any(word in prompt.lower() for word in ['population', 'populous', 'big city', 'large city', 'small city']):
    # Try direct SQL approach first
    result = city_query_engine.process_population_query(prompt)
    if result:
        message_placeholder.markdown(f"{result}\n\n*Source: Database (Direct SQL)*")
    else:
        # Fall back to LLM-based SQL
        response = sql_query_engine.query(prompt)
        message_placeholder.markdown(f"{str(response)}\n\n*Source: Database*")
elif have_llamacloud:
    # For general information, use LlamaCloud
    response = vector_query_engine.query(prompt)
    message_placeholder.markdown(f"{str(response)}\n\n*Source: LlamaCloud*")
else:
    # If neither available
    message_placeholder.markdown("I'm unable to answer that question with the current configuration.")

Error Handling and Edge Cases

The application includes several error handling mechanisms:

1. API Key Validation:

   • Checks if Google API key is present

   • Displays warning if missing

2. LlamaCloud Connection:

   • Try/except block for connection attempts

   • Fallback to SQL-only mode if connection fails

   • UI indicators for connection status

3. Query Processing:

   • Handles failed queries gracefully

   • Provides feedback on processing errors

   • Suggests query reformulation

4. Empty Results:

   • Handles case where no cities match query criteria

   • Returns informative message

Installation and Dependencies

Required Packages

streamlit==1.31.0
llama-index==0.10.0
llama-index-llms-gemini==0.1.3
llama-index-embeddings-gemini==0.1.3
llama-index-indices-managed-llama-cloud==0.1.0
llama-index-core==0.10.0
sqlalchemy==2.0.27
pandas==2.1.4

Environment Setup

1. Virtual Environment:

    python -m venv venv
    source venv/bin/activate  # On Windows: venv\Scripts\activate
   

2. Package Installation:

    pip install -r requirements.txt
   

3. API Keys:

   • Google API key for Gemini

   • LlamaCloud API key, organization ID, project, and index name

Running the Application

1. Start Application:

    streamlit run app.py
   

2. Configuration:

   • Enter API keys in sidebar

   • Confirm LlamaCloud connection status

3. Usage:

   • Ask questions in the chat input

   • View responses with source attribution

   • Chat history is maintained during session

Example Queries and Responses

Structured Data Queries

1. Query: "What is the population of New York City?" Response: "New York City has a population of 8,336,000 people and is located in New York." Source: Database (Direct SQL)

2. Query: "Which city has the highest population?" Response: "New York City has a population of 8,336,000 people and is located in New York." Source: Database (Direct SQL)

3. Query: "List all cities in California" Response: "City information:

   • Los Angeles: 3,822,000 people in California" Source: Database (Direct SQL)

Unstructured Data Queries

1. Query: "Tell me about the history of Chicago" Response: [Detailed information about Chicago's history, retrieved from documents] Source: LlamaCloud

2. Query: "What are the main attractions in Miami?" Response: [Information about Miami attractions, retrieved from documents] Source: LlamaCloud

Advanced Customization

Adding More Cities

To expand the database with additional cities:

# Add new city data
additional_rows = [
    {"city_name": "San Francisco", "population": 874961, "state": "California"},
    {"city_name": "Boston", "population": 675647, "state": "Massachusetts"},
    # Add more cities as needed
]

# Insert into database
for row in additional_rows:
    stmt = insert(city_stats_table).values(**row)
    with engine.begin() as connection:
        connection.execute(stmt)

Supporting Additional Query Types

To enhance query capabilities, extend the CityQueryEngine class:

def query_by_population_range(self, min_pop, max_pop):
    """Query cities within a population range"""
    query = f"SELECT city_name, population, state FROM city_stats WHERE population BETWEEN {min_pop} AND {max_pop} ORDER BY population DESC"
    return self.execute_query(query)

def process_population_query(self, query_text):
    # Add pattern matching for population range
    range_match = re.search(r"between\s+(\d[\d,]*)\s+and\s+(\d[\d,]*)", query_lower)
    if range_match:
        min_pop = int(range_match.group(1).replace(',', ''))
        max_pop = int(range_match.group(2).replace(',', ''))
        return self.query_by_population_range(min_pop, max_pop)

Enhancing LlamaCloud Document Retrieval

To improve document retrieval settings:

vector_query_engine = index.as_query_engine(
    similarity_top_k=3,  # Retrieve more documents
    response_mode="tree_summarize",  # Use tree summarization for better responses
    streaming=True  # Enable streaming responses
)

Troubleshooting

Common Issues and Solutions

1. API Key Errors:

   • Error: "Error: Invalid API key"

   • Solution: Verify Google API key in the sidebar or environment variables

2. LlamaCloud Connection Failures:

   • Error: "Error connecting to LlamaCloud"

   • Solution: Check organization ID, project name, and index name

3. SQL Engine Errors:

   • Error: "Error setting up SQL query engine"

   • Solution: Ensure SQLAlchemy is properly installed and database setup is correct

4. No Response to Queries:

   • Issue: System doesn't respond to certain questions

   • Solution: Check query routing logic and ensure appropriate engine is available

Performance Optimization

1. Query Caching:

   • Implement Streamlit caching for repeated queries

   • Add @st.cache_data decorator to query functions

2. Model Temperature:

   • Lower temperature (currently 0.2) for more deterministic responses

   • Increase for more creative but potentially less accurate answers

3. SQL Query Optimization:

   • Use prepared statements for frequent queries

   • Add indexes for larger datasets

Security Considerations

1. API Key Management:

   • Use environment variables instead of session state for production

   • Implement key rotation and expiration policies

2. Input Validation:

   • Sanitize user input before processing

   • Validate SQL queries to prevent injection

3. Data Privacy:

   • Consider data retention policies for chat history

   • Implement user authentication for sensitive data

Future Enhancements

1. Expanded Database:

   • Add more cities and additional attributes (GDP, crime rate, etc.)

   • Support for historical population data

2. Advanced NLP:

   • Implement intent classification for better query routing

   • Add support for multi-turn conversations with context awareness

3. UI Improvements:

   • Add visualization for city comparison

   • Implement map view for geographical context

4. Integration Options:

   • Support for additional vector databases (Qdrant, Pinecone)

   • Option to use alternative LLMs (locally via Ollama or through APIs)

5. Evaluation Framework:

   • Implement CometML's Opik for tracing and observability

   • Add metrics for answer quality and retrieval relevance

Conclusion

This LlamaCloud RAG Demo showcases an effective approach to combining structured and unstructured data for comprehensive question answering. The hybrid architecture leverages the strengths of both SQL databases (for precise factual queries) and document retrieval (for nuanced, contextual information).

Key takeaways from this implementation:

1. Intelligent Query Routing: The system automatically determines the best source for answering each question

2. Streamlined User Experience: Simple chat interface hides the complexity of the underlying systems

3. Extensible Architecture: Can be expanded with additional data sources and query capabilities

4. Production-Ready Components: Uses industry-standard tools like LlamaIndex, SQLAlchemy, and Gemini

By effectively combining these technologies, the application demonstrates a powerful approach to building knowledge-intensive applications that can handle both structured and unstructured information within a unified interface.

Google’s Gemma 3 marks a pivotal advancement in open-weight multimodal AI, demonstrating unparalleled efficiency, scalability, and safety. With enhanced long-context processing, optimized transformer architectures, and robust multimodal integration, it extends the possibilities of open-source AI. Here’s an in-depth exploration of the research innovations driving Gemma 3.

1. Pre-training: Scaling Knowledge While Ensuring Diversity

Dataset Scale and Composition

🔹 14 Trillion Tokens for the 27B Model – The sheer scale of Gemma 3’s training dataset ensures a rich and nuanced understanding of language. The focus is not just on volume but on curating a balanced dataset that generalizes well across domains.
🔹 Multilingual Expansion – Covering 140+ languages, Gemma 3 incorporates a diverse set of linguistic structures, making it one of the most comprehensive open-weight multilingual models available.
🔹 Multimodal Integration – Unlike previous iterations, Gemma 3’s pre-training includes image data, which enhances its ability to understand and generate cross-modal outputs.
🔹 Decontamination and Responsible Data Curation – A rigorous decontamination process was employed to mitigate risks related to memorization, sensitive data leakage, and bias, reinforcing Google’s commitment to responsible AI development.

Tokenization: Optimizing for Efficiency

🔹 SentencePiece Tokenizer – Chosen for its subword tokenization capabilities, SentencePiece allows Gemma 3 to effectively process out-of-vocabulary words, morphologically rich languages, and diverse scripts.
🔹 Efficiency Gains – By refining tokenization strategies, the model achieves better compression and improved language modeling, reducing unnecessary token overhead.

Computational Infrastructure: Scaling Training Efficiently

🔹 Trained on Google TPUs (TPUv4p, TPUv5p, TPUv5e) – These next-generation accelerators provide unmatched efficiency in matrix computations, reducing both training time and energy consumption.
🔹 Optimized for Large-Scale Training – Leveraging distributed training strategies, Gemma 3 maximizes hardware utilization, ensuring the best performance-to-compute ratio.

# SECTION 1: Pre-training - Dataset Processing and Tokenization
"""
This example demonstrates how to process a multilingual dataset and apply
tokenization techniques similar to those used in Gemma 3's pre-training.
"""

import numpy as np
import pandas as pd
from typing import List, Dict, Any
from sentencepiece import SentencePieceProcessor

class GemmaDatasetProcessor:
    def __init__(self, tokenizer_path: str, languages: List[str] = None):
        """Initialize the dataset processor with a SentencePiece tokenizer."""
        self.tokenizer = SentencePieceProcessor()
        self.tokenizer.Load(tokenizer_path)
        self.languages = languages or ["en", "fr", "de", "es", "zh", "ja", "ar", "hi"]
        self.stats = {"token_count": 0, "examples_per_lang": {lang: 0 for lang in self.languages}}

    def load_and_process_data(self, data_paths: Dict[str, str]) -> List[Dict[str, Any]]:
        """Load and process data from multiple sources, keeping track of statistics."""
        processed_data = []

        for lang, path in data_paths.items():
            if lang not in self.languages:
                continue

            print(f"Processing {lang} data from {path}")
            # Load data (simplified example)
            raw_texts = pd.read_csv(path)["text"].tolist()

            for text in raw_texts:
                # Apply tokenization
                tokens = self.tokenizer.EncodeAsIds(text)

                # Apply decontamination logic (simplified)
                is_clean = self._decontaminate_text(text)

                if is_clean and len(tokens) > 0:
                    processed_example = {
                        "language": lang,
                        "text": text,
                        "tokens": tokens,
                        "n_tokens": len(tokens)
                    }
                    processed_data.append(processed_example)

                    # Update statistics
                    self.stats["token_count"] += len(tokens)
                    self.stats["examples_per_lang"][lang] += 1

        print(f"Processed {len(processed_data)} examples with {self.stats['token_count']} tokens")
        return processed_data

    def _decontaminate_text(self, text: str) -> bool:
        """
        Simplified decontamination that checks for sensitive patterns.
        In a real implementation, this would be much more sophisticated.
        """
        sensitive_patterns = ["password:", "secret:", "private key"]
        return not any(pattern in text.lower() for pattern in sensitive_patterns)

    def get_dataset_stats(self) -> Dict[str, Any]:
        """Return statistics about the processed dataset."""
        return {
            "total_tokens": self.stats["token_count"],
            "languages": len(self.languages),
            "language_distribution": {
                lang: count / sum(self.stats["examples_per_lang"].values())
                for lang, count in self.stats["examples_per_lang"].items() if count > 0
            }
        }

# Example usage (not executed)
if __name__ == "__main__":
    processor = GemmaDatasetProcessor("path/to/sentencepiece_model.model")
    data_paths = {
        "en": "path/to/english_data.csv",
        "fr": "path/to/french_data.csv",
        # Add more languages
    }
    processed_data = processor.load_and_process_data(data_paths)
    stats = processor.get_dataset_stats()
    print(f"Dataset statistics: {stats}")

2. Architectural Innovations: Balancing Efficiency and Scalability

Enhanced Attention Mechanisms

🔹 Interleaved Local and Global Attention – This novel strategy mitigates the quadratic complexity of self-attention, allowing Gemma 3 to efficiently process longer contexts (up to 128K tokens).
🔹 Grouped-Query Attention (GQA) & QK-Norm – These refinements reduce memory overhead while maintaining high-quality representations, making the model more practical for real-time and production applications.
🔹 Refined Rotary Positional Embeddings (RoPE) – Adjustments to RoPE base frequency improve the model’s ability to track relationships over long sequences, further strengthening its context retention.

Multimodal Capabilities: A Step Towards General Intelligence

🔹 SigLIP Vision Encoder – A key addition enabling seamless image-text processing, making the model more adept at handling multimodal queries.
🔹 Adaptive Windowing for Images – This technique allows Gemma 3 to dynamically process images of varying resolutions and aspect ratios, improving its visual comprehension abilities.

# SECTION 2: Architectural Innovations - Implementing Enhanced Attention Mechanisms
"""
This example demonstrates how to implement the enhanced attention mechanisms 
used in Gemma 3, including interleaved local and global attention and grouped-query attention.
"""

import torch
import torch.nn as nn
import math
from typing import Optional, Tuple

class GemmaRotaryEmbedding(nn.Module):
    """Rotary positional embeddings with frequency adjustments."""

    def __init__(self, dim: int, base=10000.0, scaling_factor=1.0):
        super().__init__()
        self.dim = dim
        self.base = base * scaling_factor
        inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, seq_len: int, device: torch.device):
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb

def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    """Apply rotary position embeddings to query and key tensors."""
    # Reshape for broadcasting
    cos = cos[:, :, None, :]  # [batch, seq_len, 1, dim]
    sin = sin[:, :, None, :]  # [batch, seq_len, 1, dim]

    # Apply rotation using complex multiplication logic
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)

    return q_embed, k_embed

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Rotate half of the hidden dims."""
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

class GemmaAttention(nn.Module):
    """
    Enhanced attention module implementing interleaved local and global attention
    with Grouped-Query Attention (GQA) and QK-Norm.
    """

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int = None,  # For GQA
        window_size: int = 1024,  # For local attention
        qk_norm: bool = True,
        dropout_prob: float = 0.0,
    ):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads or num_heads
        self.head_dim = hidden_size // num_heads
        self.window_size = window_size
        self.qk_norm = qk_norm

        # Calculate grouping factor for GQA
        self.num_groups = self.num_heads // self.num_kv_heads

        # Projection matrices
        self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)

        # Rotary embeddings
        self.rotary_emb = GemmaRotaryEmbedding(self.head_dim, scaling_factor=0.1)  # Adjusted base freq

        # Layer norms for QK-Norm
        if qk_norm:
            self.q_norm = nn.LayerNorm(self.head_dim)
            self.k_norm = nn.LayerNorm(self.head_dim)

        self.dropout = nn.Dropout(dropout_prob)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        is_local: bool = True,  # Switch between local and global attention
        position_ids: Optional[torch.LongTensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass with support for interleaved local and global attention.

        Args:
            hidden_states: Input tensor [batch_size, seq_len, hidden_size]
            attention_mask: Attention mask [batch_size, 1, 1, seq_len]
            is_local: Whether to use local attention (True) or global attention (False)
            position_ids: Optional explicit position IDs

        Returns:
            output: Output tensor [batch_size, seq_len, hidden_size]
            attention_weights: Attention weights for visualization
        """
        batch_size, seq_length, _ = hidden_states.shape

        # Project inputs to queries, keys, and values
        q = self.q_proj(hidden_states).view(batch_size, seq_length, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)

        # Apply QK-Norm if enabled
        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)

        # Apply rotary embeddings
        if position_ids is None:
            position_ids = torch.arange(seq_length, device=hidden_states.device)

        cos, sin = self.rotary_emb(seq_length, hidden_states.device)
        cos = cos[position_ids].unsqueeze(0)  # [1, seq_len, dim]
        sin = sin[position_ids].unsqueeze(0)  # [1, seq_len, dim]

        q, k = apply_rotary_pos_emb(q, k, cos, sin)

        # Grouped-Query Attention: Repeat KV heads for each query group
        if self.num_groups > 1:
            k = k.unsqueeze(2).expand(-1, -1, self.num_groups, -1, -1)
            v = v.unsqueeze(2).expand(-1, -1, self.num_groups, -1, -1)
            k = k.reshape(batch_size, seq_length, self.num_heads, self.head_dim)
            v = v.reshape(batch_size, seq_length, self.num_heads, self.head_dim)

        # Transpose for attention computation
        q = q.transpose(1, 2)  # [batch_size, num_heads, seq_length, head_dim]
        k = k.transpose(1, 2)  # [batch_size, num_heads, seq_length, head_dim]
        v = v.transpose(1, 2)  # [batch_size, num_heads, seq_length, head_dim]

        # Create local or global attention pattern
        if is_local:
            # Local attention with window
            attention_weights = torch.zeros(
                batch_size, self.num_heads, seq_length, seq_length, 
                device=hidden_states.device
            )

            # For each position, attend only to nearby positions within window
            for i in range(seq_length):
                start = max(0, i - self.window_size // 2)
                end = min(seq_length, i + self.window_size // 2 + 1)

                local_q = q[:, :, i:i+1]  # [batch_size, num_heads, 1, head_dim]
                local_k = k[:, :, start:end]  # [batch_size, num_heads, window_size, head_dim]
                local_v = v[:, :, start:end]  # [batch_size, num_heads, window_size, head_dim]

                # Compute attention scores for this position
                attn_scores = torch.matmul(local_q, local_k.transpose(2, 3)) / math.sqrt(self.head_dim)

                if attention_mask is not None:
                    local_mask = attention_mask[:, :, :, start:end]
                    attn_scores = attn_scores + local_mask

                attn_probs = torch.softmax(attn_scores, dim=-1)
                attn_probs = self.dropout(attn_probs)

                # Update the value for this position
                local_output = torch.matmul(attn_probs, local_v)
                attention_weights[:, :, i, start:end] = attn_probs.squeeze(2)

                if i == 0:
                    output = local_output
                else:
                    output = torch.cat([output, local_output], dim=2)
        else:
            # Global attention (standard self-attention)
            attention_scores = torch.matmul(q, k.transpose(2, 3)) / math.sqrt(self.head_dim)

            if attention_mask is not None:
                attention_scores = attention_scores + attention_mask

            attention_weights = torch.softmax(attention_scores, dim=-1)
            attention_weights = self.dropout(attention_weights)

            output = torch.matmul(attention_weights, v)

        # Restore original shape
        output = output.transpose(1, 2).reshape(batch_size, seq_length, self.hidden_size)

        # Final projection
        output = self.o_proj(output)

        return output, attention_weights

3. Post-training: Refining the Model Through Alignment and Safety

Knowledge Distillation for Efficient Deployment

🔹 Distillation from Large Models to Smaller Variants – By leveraging teacher-student learning, smaller models inherit the capabilities of larger counterparts, making them ideal for resource-constrained environments (e.g., edge devices).

Instruction Tuning (IT) and Reinforcement Learning (RL)

🔹 Improved Instruction Following – IT significantly enhances Gemma 3’s ability to understand, execute, and generalize user instructions.
🔹 RLHF & RLMF for Alignment – Reinforcement Learning from Human Feedback (RLHF) and Machine Feedback (RLMF) optimize the model’s reasoning skills, particularly improving its mathematical and logical consistency.

Safety and Responsible AI

🔹 Robust Safety Evaluations – The model underwent extensive red-teaming and adversarial testing to detect and mitigate harmful content generation, biases, and privacy risks.
🔹 Alignment for Ethical AI – Special attention was given to reducing toxic outputs, representational harms, and information leakage, ensuring Gemma 3 meets high ethical and safety standards.

# SECTION 3: Post-training - Knowledge Distillation and Safety
"""
This example demonstrates how to implement knowledge distillation to transfer knowledge
from a large teacher model to a smaller student model, as well as a basic safety classifier.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Dict, List, Tuple

class KnowledgeDistillationTrainer:
    """Trainer for knowledge distillation from a larger Gemma 3 model to a smaller one."""

    def __init__(
        self,
        teacher_model_id: str,
        student_model_id: str,
        alpha: float = 0.5,  # Balance between distillation and ground truth
        temperature: float = 2.0,  # Temperature for softening probability distributions
    ):
        # Initialize teacher model (larger model like Gemma 3 27B)
        self.teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model_id)
        self.teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model_id)

        # Initialize student model (smaller model to be trained)
        self.student_model = AutoModelForCausalLM.from_pretrained(student_model_id)
        self.student_tokenizer = AutoTokenizer.from_pretrained(student_model_id)

        # Freeze teacher parameters
        for param in self.teacher_model.parameters():
            param.requires_grad = False

        self.alpha = alpha
        self.temperature = temperature

        # Optimizer for student model
        self.optimizer = torch.optim.AdamW(self.student_model.parameters(), lr=5e-5)

    def compute_distillation_loss(
        self,
        teacher_logits: torch.Tensor,
        student_logits: torch.Tensor,
        target_ids: torch.Tensor,
        attention_mask: torch.Tensor
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """
        Compute the knowledge distillation loss.

        Args:
            teacher_logits: Logits from teacher model [batch, seq_len, vocab_size]
            student_logits: Logits from student model [batch, seq_len, vocab_size]
            target_ids: Target token IDs [batch, seq_len]
            attention_mask: Attention mask [batch, seq_len]

        Returns:
            total_loss: Combined distillation and cross-entropy loss
            loss_dict: Dictionary containing individual loss components
        """
        # Apply temperature to soften probability distributions
        soft_teacher_logits = teacher_logits / self.temperature
        soft_student_logits = student_logits / self.temperature

        # Compute KL divergence loss for knowledge distillation
        distillation_loss = F.kl_div(
            F.log_softmax(soft_student_logits, dim=-1),
            F.softmax(soft_teacher_logits, dim=-1),
            reduction="batchmean"
        ) * (self.temperature ** 2)

        # Compute cross-entropy loss against ground truth
        ce_loss = F.cross_entropy(
            student_logits.view(-1, student_logits.size(-1)),
            target_ids.view(-1),
            ignore_index=-100
        )

        # Combine losses
        total_loss = self.alpha * distillation_loss + (1 - self.alpha) * ce_loss

        return total_loss, {
            "total_loss": total_loss.item(),
            "distillation_loss": distillation_loss.item(),
            "ce_loss": ce_loss.item()
        }

    def train_step(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
        """Perform a single training step with knowledge distillation."""
        # Move batch to device
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        target_ids = batch["labels"]

        # Forward pass through teacher model (no gradients)
        with torch.no_grad():
            teacher_outputs = self.teacher_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_dict=True
            )
            teacher_logits = teacher_outputs.logits

        # Forward pass through student model
        student_outputs = self.student_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )
        student_logits = student_outputs.logits

        # Compute loss
        loss, loss_dict = self.compute_distillation_loss(
            teacher_logits, student_logits, target_ids, attention_mask
        )

        # Backward pass and update student parameters
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return loss_dict


class SafetyClassifier(nn.Module):
    """Safety classifier based on a pre-trained model to detect harmful content."""

    def __init__(self, base_model_id: str, num_safety_categories: int = 8):
        super().__init__()
        # Load base model
        self.tokenizer = AutoTokenizer.from_pretrained(base_model_id)
        self.encoder = AutoModelForCausalLM.from_pretrained(base_model_id)

        # Freeze base model
        for param in self.encoder.parameters():
            param.requires_grad = False

        # Classification head
        self.safety_head = nn.Sequential(
            nn.Linear(self.encoder.config.hidden_size, 512),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(512, num_safety_categories)
        )

        # Safety categories
        self.safety_categories = [
            "hate_speech",
            "harassment",
            "self_harm",
            "sexual_content",
            "violence",
            "dangerous_content",
            "misinformation",
            "personal_data"
        ]

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Forward pass to classify text for safety risks."""
        # Get embeddings from base model
        with torch.no_grad():
            outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=True,
                return_dict=True
            )

            # Use the last hidden state of the final token
            last_token_indices = attention_mask.sum(dim=1) - 1
            batch_indices = torch.arange(input_ids.size(0))
            last_token_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices]

        # Pass through safety classification head
        safety_logits = self.safety_head(last_token_hidden)
        return safety_logits

    def classify_text(self, text: str, threshold: float = 0.5) -> Dict[str, float]:
        """Classify a text for safety risks."""
        # Tokenize text
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=512,
            padding=True
        )

        # Get safety logits
        safety_logits = self(inputs.input_ids, inputs.attention_mask)

        # Apply sigmoid to get probabilities
        safety_probs = torch.sigmoid(safety_logits).squeeze().detach().numpy()

        # Create results dictionary
        results = {
            category: float(prob)
            for category, prob in zip(self.safety_categories, safety_probs)
        }

        # Add overall safety assessment
        results["is_unsafe"] = any(prob > threshold for prob in safety_probs)
        results["max_risk_category"] = self.safety_categories[safety_probs.argmax()]
        results["max_risk_score"] = float(safety_probs.max())

        return results

4. Benchmarking: Performance at Scale

Competitive Performance Against Leading Open Models

✅ Outperforms Llama 3 and DeepSeek R1 across multiple language and reasoning benchmarks.
✅ Excels in multilingual NLP tasks, leveraging its diverse training data.

Long-Context Mastery

✅ With a 128K token context window, Gemma 3 can retain and utilize information over extended sequences, making it ideal for applications requiring document synthesis, code completion, and legal/academic analysis.

Safety and Trustworthiness

✅ Significant improvements in safety benchmarks, reflecting a lower propensity for harmful or misleading outputs.

# SECTION 3: Post-training - Knowledge Distillation and Safety
"""
This example demonstrates how to implement knowledge distillation to transfer knowledge
from a large teacher model to a smaller student model, as well as a basic safety classifier.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Dict, List, Tuple

class KnowledgeDistillationTrainer:
    """Trainer for knowledge distillation from a larger Gemma 3 model to a smaller one."""

    def __init__(
        self,
        teacher_model_id: str,
        student_model_id: str,
        alpha: float = 0.5,  # Balance between distillation and ground truth
        temperature: float = 2.0,  # Temperature for softening probability distributions
    ):
        # Initialize teacher model (larger model like Gemma 3 27B)
        self.teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model_id)
        self.teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_model_id)

        # Initialize student model (smaller model to be trained)
        self.student_model = AutoModelForCausalLM.from_pretrained(student_model_id)
        self.student_tokenizer = AutoTokenizer.from_pretrained(student_model_id)

        # Freeze teacher parameters
        for param in self.teacher_model.parameters():
            param.requires_grad = False

        self.alpha = alpha
        self.temperature = temperature

        # Optimizer for student model
        self.optimizer = torch.optim.AdamW(self.student_model.parameters(), lr=5e-5)

    def compute_distillation_loss(
        self,
        teacher_logits: torch.Tensor,
        student_logits: torch.Tensor,
        target_ids: torch.Tensor,
        attention_mask: torch.Tensor
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """
        Compute the knowledge distillation loss.

        Args:
            teacher_logits: Logits from teacher model [batch, seq_len, vocab_size]
            student_logits: Logits from student model [batch, seq_len, vocab_size]
            target_ids: Target token IDs [batch, seq_len]
            attention_mask: Attention mask [batch, seq_len]

        Returns:
            total_loss: Combined distillation and cross-entropy loss
            loss_dict: Dictionary containing individual loss components
        """
        # Apply temperature to soften probability distributions
        soft_teacher_logits = teacher_logits / self.temperature
        soft_student_logits = student_logits / self.temperature

        # Compute KL divergence loss for knowledge distillation
        distillation_loss = F.kl_div(
            F.log_softmax(soft_student_logits, dim=-1),
            F.softmax(soft_teacher_logits, dim=-1),
            reduction="batchmean"
        ) * (self.temperature ** 2)

        # Compute cross-entropy loss against ground truth
        ce_loss = F.cross_entropy(
            student_logits.view(-1, student_logits.size(-1)),
            target_ids.view(-1),
            ignore_index=-100
        )

        # Combine losses
        total_loss = self.alpha * distillation_loss + (1 - self.alpha) * ce_loss

        return total_loss, {
            "total_loss": total_loss.item(),
            "distillation_loss": distillation_loss.item(),
            "ce_loss": ce_loss.item()
        }

    def train_step(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]:
        """Perform a single training step with knowledge distillation."""
        # Move batch to device
        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        target_ids = batch["labels"]

        # Forward pass through teacher model (no gradients)
        with torch.no_grad():
            teacher_outputs = self.teacher_model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_dict=True
            )
            teacher_logits = teacher_outputs.logits

        # Forward pass through student model
        student_outputs = self.student_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True
        )
        student_logits = student_outputs.logits

        # Compute loss
        loss, loss_dict = self.compute_distillation_loss(
            teacher_logits, student_logits, target_ids, attention_mask
        )

        # Backward pass and update student parameters
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return loss_dict


class SafetyClassifier(nn.Module):
    """Safety classifier based on a pre-trained model to detect harmful content."""

    def __init__(self, base_model_id: str, num_safety_categories: int = 8):
        super().__init__()
        # Load base model
        self.tokenizer = AutoTokenizer.from_pretrained(base_model_id)
        self.encoder = AutoModelForCausalLM.from_pretrained(base_model_id)

        # Freeze base model
        for param in self.encoder.parameters():
            param.requires_grad = False

        # Classification head
        self.safety_head = nn.Sequential(
            nn.Linear(self.encoder.config.hidden_size, 512),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(512, num_safety_categories)
        )

        # Safety categories
        self.safety_categories = [
            "hate_speech",
            "harassment",
            "self_harm",
            "sexual_content",
            "violence",
            "dangerous_content",
            "misinformation",
            "personal_data"
        ]

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        """Forward pass to classify text for safety risks."""
        # Get embeddings from base model
        with torch.no_grad():
            outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=True,
                return_dict=True
            )

            # Use the last hidden state of the final token
            last_token_indices = attention_mask.sum(dim=1) - 1
            batch_indices = torch.arange(input_ids.size(0))
            last_token_hidden = outputs.hidden_states[-1][batch_indices, last_token_indices]

        # Pass through safety classification head
        safety_logits = self.safety_head(last_token_hidden)
        return safety_logits

    def classify_text(self, text: str, threshold: float = 0.5) -> Dict[str, float]:
        """Classify a text for safety risks."""
        # Tokenize text
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=512,
            padding=True
        )

        # Get safety logits
        safety_logits = self(inputs.input_ids, inputs.attention_mask)

        # Apply sigmoid to get probabilities
        safety_probs = torch.sigmoid(safety_logits).squeeze().detach().numpy()

        # Create results dictionary
        results = {
            category: float(prob)
            for category, prob in zip(self.safety_categories, safety_probs)
        }

        # Add overall safety assessment
        results["is_unsafe"] = any(prob > threshold for prob in safety_probs)
        results["max_risk_category"] = self.safety_categories[safety_probs.argmax()]
        results["max_risk_score"] = float(safety_probs.max())

        return results

Final Thoughts: The Future of Open AI

Gemma 3 represents a significant advancement in scalable, efficient, and safe AI, demonstrating Google’s commitment to open research and responsible AI development.

With its cutting-edge transformer refinements, long-context capabilities, and multimodal integration, this model is poised to accelerate progress in AI research and real-world applications.

Gemma 3 Blog - Google

Gemma 3 Technical Report

Conclusion

Gemma 3 stands as a remarkable leap forward in the realm of open-source multimodal AI, showcasing Google's dedication to advancing AI technology responsibly and efficiently. With its innovative transformer architectures, extensive multilingual capabilities, and robust multimodal integration, Gemma 3 is set to drive significant progress in AI research and practical applications. Its ability to process long contexts and handle diverse tasks with improved safety and ethical standards positions it as a pivotal tool for future developments in AI. As the landscape of AI continues to evolve, Gemma 3 exemplifies the potential of open-weight models to contribute to a more inclusive and advanced technological future.

#AI #MachineLearning #LLM #OpenSource #MultimodalAI #Gemma3 #GoogleAI #GoogleDeepMind