AIModelHandler

​

Overview

​

Class: AIModelHandler

class AIModelHandler:
    """Handles AI model operations including initialization and text generation."""

​

Attributes

• models: Dict[str, Any] - Dictionary storing loaded model instances
• tokenizers: Dict[str, Any] - Dictionary storing tokenizer instances
• model_configs: Dict[str, Dict] - Configuration for each supported model
• default_model: str - Default model key (from DEFAULT_MODEL env var, defaults to ‘llama’)
• enable_learning: bool - Continuous learning flag (from ENABLE_CONTINUOUS_LEARNING env var)

​

Methods

​

__init__()

def __init__(self)

'llama': {
    'name': 'nvidia/Llama-3.1-Nemotron-70B-Instruct-HF',
    'type': 'causal',
    'task': 'text-generation'
}

'bart': {
    'name': 'facebook/bart-large',
    'type': 'conditional',
    'task': 'summarization'
}

• DEFAULT_MODEL: Model key to use by default (default: ‘llama’)
• ENABLE_CONTINUOUS_LEARNING: Enable learning features (default: ‘true’)

​

initialize_models()

async def initialize_models(self) -> bool

1. GPU Detection
   • Checks CUDA availability
   • Logs GPU device name and available memory
2. BART Model Initialization
   • Loads facebook/bart-large tokenizer
   • Loads BART model with automatic device mapping
   • Uses float16 precision on GPU, float32 on CPU
3. Llama Model Initialization
   • Loads nvidia/Llama-3.1-Nemotron-70B-Instruct-HF tokenizer
   • Loads Llama model with automatic device mapping
   • Uses float16 precision on GPU, float32 on CPU

• device_map='auto': Automatic device placement when GPU is available
• torch_dtype: float16 for GPU (memory efficient), float32 for CPU
• local_files_only=False: Downloads models from Hugging Face if not cached

• Returns False if any model fails to load
• Logs detailed error messages for debugging
• Continues initialization if one model fails (when possible)

​

generate_response()

async def generate_response(
    self,
    text: str,
    model_key: Optional[str] = None,
    max_length: int = 300,
    temperature: float = 0.2,
    top_p: float = 0.4,
    max_attempts: int = 5
) -> str

1. Model Selection
   • Uses specified model_key or falls back to default_model
   • Validates that the model exists
2. Input Processing
   • Tokenizes input text
   • Truncates to 512 tokens maximum
   • Adds padding for batch processing
   • Moves tensors to model’s device (CPU/GPU)
3. Generation Loop
   • Attempts generation up to max_attempts times
   • Uses nucleus sampling with specified temperature and top_p
   • Validates response quality (>3 words, different from input)
   • Breaks on first valid response
4. Response Validation
   • Checks response length (must be >3 words)
   • Ensures response differs from input
   • Returns fallback message if all attempts fail

• do_sample=True: Enables sampling for diverse outputs
• num_return_sequences=1: Generates one response
• Uses model’s pad_token_id and eos_token_id for proper sequence handling

• Raises ValueError if model_key is invalid
• Returns user-friendly error message on exceptions
• Logs all errors for debugging

​

summarize_text()

async def summarize_text(
    self,
    text: str,
    max_length: int = 130,
    min_length: int = 30
) -> str

1. Input Processing
   • Tokenizes input text using BART tokenizer
   • Truncates to 1024 tokens maximum
   • Adds padding for batch processing
   • Moves tensors to BART model’s device
2. Summary Generation
   • Uses beam search with 4 beams for better quality
   • Applies length penalty of 2.0 to encourage conciseness
   • Early stopping when all beams reach EOS token
   • Decodes output tokens to text
3. Post-processing
   • Removes special tokens from output
   • Strips whitespace from final summary

• num_beams=4: Beam search width for quality
• length_penalty=2.0: Encourages shorter summaries
• early_stopping=True: Stops when EOS is reached

• Returns user-friendly error message on exceptions
• Logs all errors for debugging

​

get_available_models()

def get_available_models(self) -> List[str]

• Returns the keys from model_configs dictionary
• Useful for validating model selection

​

get_model_info()

def get_model_info(self, model_key: str) -> Optional[Dict]

{
    'name': 'nvidia/Llama-3.1-Nemotron-70B-Instruct-HF',
    'type': 'causal',
    'task': 'text-generation'
}

​

Singleton Instance

ai_handler = AIModelHandler()

​

Dependencies

import logging
import os
from typing import Dict, Optional, List, Any
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BartForConditionalGeneration,
    BartTokenizer,
    pipeline
)

​

Environment Variables

• DEFAULT_MODEL (optional): Model key to use by default (default: ‘llama’)
• ENABLE_CONTINUOUS_LEARNING (optional): Enable continuous learning features (default: ‘true’)

​

GPU Support

• Uses CUDA if torch.cuda.is_available() returns True
• Automatically maps models to available devices
• Uses float16 precision on GPU for memory efficiency
• Falls back to CPU with float32 precision when GPU is unavailable