from langfuse import Langfuse from langfuse.decorators import observe, langfuse_context from config.config import settings import os # Initialize Langfuse os.environ["LANGFUSE_PUBLIC_KEY"] = "pk-lf-9f2c32d2-266f-421d-9b87-51377f0a268c" os.environ["LANGFUSE_SECRET_KEY"] = "sk-lf-229e10c5-6210-4a4b-a432-0f17bc66e56c" os.environ["LANGFUSE_HOST"] = "https://chris4k-langfuse-template-space.hf.space" # 🇪🇺 EU region try: langfuse = Langfuse() except Exception as e: print("Langfuse Offline") # model_manager.py import torch from transformers import AutoModelForCausalLM, AutoTokenizer from llama_cpp import Llama from typing import Optional, Dict import logging from functools import lru_cache from config.config import GenerationConfig, ModelConfig class ModelManager: def __init__(self, device: Optional[str] = None): self.logger = logging.getLogger(__name__) self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.models: Dict[str, Any] = {} self.tokenizers: Dict[str, Any] = {} @observe() def load_model(self, model_id: str, model_path: str, model_type: str, config: ModelConfig) -> None: """Load a model with specified configuration.""" try: ##could be differnt models, so we can use a factory pattern to load the correct model - textgen, llama, gguf, text2video, text2image etc. if model_type == "llama": self.tokenizers[model_id] = AutoTokenizer.from_pretrained( model_path, padding_side='left', trust_remote_code=True, **config.tokenizer_kwargs ) if self.tokenizers[model_id].pad_token is None: self.tokenizers[model_id].pad_token = self.tokenizers[model_id].eos_token self.models[model_id] = AutoModelForCausalLM.from_pretrained( model_path, device_map="auto", trust_remote_code=True, **config.model_kwargs ) elif model_type == "gguf": #TODO load the model first from the cache, if not found load the model and save it in the cache #from huggingface_hub import hf_hub_download #prm_model_path = hf_hub_download( # repo_id="tensorblock/Llama3.1-8B-PRM-Mistral-Data-GGUF", # filename="Llama3.1-8B-PRM-Mistral-Data-Q4_K_M.gguf" #) self.models[model_id] = self._load_quantized_model( model_path, **config.quantization_kwargs ) except Exception as e: self.logger.error(f"Failed to load model {model_id}: {str(e)}") raise @observe() def unload_model(self, model_id: str) -> None: """Unload a model and free resources.""" if model_id in self.models: del self.models[model_id] if model_id in self.tokenizers: del self.tokenizers[model_id] torch.cuda.empty_cache() def _load_quantized_model(self, model_path: str, **kwargs) -> Llama: """Load a quantized GGUF model.""" try: n_gpu_layers = -1 if torch.cuda.is_available() else 0 model = Llama( model_path=model_path, n_ctx=kwargs.get('n_ctx', 2048), n_batch=kwargs.get('n_batch', 512), n_gpu_layers=kwargs.get('n_gpu_layers', n_gpu_layers), verbose=kwargs.get('verbose', False) ) return model except Exception as e: self.logger.error(f"Failed to load GGUF model: {str(e)}") raise # cache.py from functools import lru_cache from typing import Tuple, Any # TODO explain howto use the cache class ResponseCache: def __init__(self, cache_size: int = 1000): self.cache_size = cache_size self._initialize_cache() def _initialize_cache(self): @lru_cache(maxsize=self.cache_size) def cached_response(prompt: str, config_hash: str) -> Tuple[str, float]: pass self.get_cached_response = cached_response def cache_response(self, prompt: str, config: GenerationConfig, response: str, score: float) -> None: config_hash = hash(str(config.__dict__)) self.get_cached_response(prompt, str(config_hash)) def get_response(self, prompt: str, config: GenerationConfig) -> Optional[Tuple[str, float]]: config_hash = hash(str(config.__dict__)) return self.get_cached_response(prompt, str(config_hash)) # batch_processor.py from typing import List, Dict import asyncio #TODO explain how to use the batch processor class BatchProcessor: def __init__(self, max_batch_size: int = 32, max_wait_time: float = 0.1): self.max_batch_size = max_batch_size self.max_wait_time = max_wait_time self.pending_requests: List[Dict] = [] self.lock = asyncio.Lock() async def add_request(self, request: Dict) -> Any: async with self.lock: self.pending_requests.append(request) if len(self.pending_requests) >= self.max_batch_size: return await self._process_batch() else: await asyncio.sleep(self.max_wait_time) if self.pending_requests: return await self._process_batch() async def _process_batch(self) -> List[Any]: batch = self.pending_requests[:self.max_batch_size] self.pending_requests = self.pending_requests[self.max_batch_size:] # TODO implement the batch processing logic return batch # base_generator.py from abc import ABC, abstractmethod from typing import AsyncGenerator, Dict, Any, Optional, List, Tuple from dataclasses import dataclass from logging import getLogger from config.config import GenerationConfig, ModelConfig class BaseGenerator(ABC): """Base class for all generator implementations.""" def __init__( self, model_name: str, device: Optional[str] = None, default_generation_config: Optional[GenerationConfig] = None, model_config: Optional[ModelConfig] = None, cache_size: int = 1000, max_batch_size: int = 32 ): self.logger = getLogger(__name__) self.model_manager = ModelManager(device) self.cache = ResponseCache(cache_size) self.batch_processor = BatchProcessor(max_batch_size) self.health_check = HealthCheck() # self.tokenizer = self.model_manager.tokenizers[model_name] #self.tokenizer = self.load_tokenizer(llama_model_name) # Add this line to initialize the tokenizer self.default_config = default_generation_config or GenerationConfig() self.model_config = model_config or ModelConfig() @abstractmethod async def generate_stream( self, prompt: str, config: Optional[GenerationConfig] = None ) -> AsyncGenerator[str, None]: pass @abstractmethod def _get_generation_kwargs(self, config: GenerationConfig) -> Dict[str, Any]: pass @abstractmethod def generate( self, prompt: str, model_kwargs: Dict[str, Any], strategy: str = "default", **kwargs ) -> str: pass # strategy.py #TODO UPDATE Paths from abc import ABC, abstractmethod from typing import List, Tuple @observe() class GenerationStrategy(ABC): """Base class for generation strategies.""" @abstractmethod def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs) -> str: pass class DefaultStrategy(GenerationStrategy): def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs) -> str: input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device) output = generator.model.generate(input_ids, **model_kwargs) return generator.tokenizer.decode(output[0], skip_special_tokens=True) @observe() class MajorityVotingStrategy(GenerationStrategy): def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str: outputs = [] for _ in range(num_samples): input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device) output = generator.model.generate(input_ids, **model_kwargs) outputs.append(generator.tokenizer.decode(output[0], skip_special_tokens=True)) return max(set(outputs), key=outputs.count) @observe() class BestOfN(GenerationStrategy): def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str: scored_outputs = [] for _ in range(num_samples): input_ids = self.llama_tokenizer(prompt, return_tensors="pt").input_ids.to(self.device) output = self.llama_model.generate(input_ids, **model_kwargs) response = self.llama_tokenizer.decode(output[0], skip_special_tokens=True) score = self.prm_model(**self.llama_tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item() scored_outputs.append((response, score)) return max(scored_outputs, key=lambda x: x[1])[0] @observe() class BeamSearch(GenerationStrategy): def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str: input_ids = self.llama_tokenizer(prompt, return_tensors="pt").input_ids.to(self.device) outputs = self.llama_model.generate( input_ids, num_beams=num_samples, num_return_sequences=num_samples, **model_kwargs ) return [self.llama_tokenizer.decode(output, skip_special_tokens=True) for output in outputs] @observe() class DVT(GenerationStrategy): def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str: results = [] for _ in range(breadth): input_ids = self.llama_tokenizer(prompt, return_tensors="pt").input_ids.to(self.device) output = self.llama_model.generate(input_ids, **model_kwargs) response = self.llama_tokenizer.decode(output[0], skip_special_tokens=True) score = self.prm_model(**self.llama_tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item() results.append((response, score)) for _ in range(depth - 1): best_responses = sorted(results, key=lambda x: x[1], reverse=True)[:breadth] for response, _ in best_responses: input_ids = self.llama_tokenizer(response, return_tensors="pt").input_ids.to(self.device) output = self.llama_model.generate(input_ids, **model_kwargs) extended_response = self.llama_tokenizer.decode(output[0], skip_special_tokens=True) score = self.prm_model(**self.llama_tokenizer(extended_response, return_tensors="pt").to(self.device)).logits.mean().item() results.append((extended_response, score)) return max(results, key=lambda x: x[1])[0] @observe() class COT(GenerationStrategy): def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str: #TODO implement the chain of thought strategy return "Not implemented yet" @observe() class ReAct(GenerationStrategy): def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str: #TODO implement the ReAct framework return "Not implemented yet" # Add other strategy implementations... # prompt_builder.py from typing import Protocol, List, Tuple from transformers import AutoTokenizer @observe() class PromptTemplate(Protocol): """Protocol for prompt templates.""" def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], **kwargs) -> str: pass @observe() class LlamaPromptTemplate: def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], max_history_turns: int = 1) -> str: system_message = f"Please assist based on the following context: {context}" prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>" for user_msg, assistant_msg in chat_history[-max_history_turns:]: prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_msg}<|eot_id|>" prompt += f"<|start_header_id|>assistant<|end_header_id|>\n\n{assistant_msg}<|eot_id|>" prompt += f"<|start_header_id|>user<|end_header_id|>\n\n{user_input}<|eot_id|>" prompt += "<|start_header_id|>assistant<|end_header_id|>\n\n" return prompt @observe() class TransformersPromptTemplate: def __init__(self, model_path: str): self.tokenizer = AutoTokenizer.from_pretrained(model_path) def format(self, context: str, user_input: str, chat_history: List[Tuple[str, str]], **kwargs) -> str: messages = [ { "role": "system", "content": f"Please assist based on the following context: {context}", } ] for user_msg, assistant_msg in chat_history: messages.extend([ {"role": "user", "content": user_msg}, {"role": "assistant", "content": assistant_msg} ]) messages.append({"role": "user", "content": user_input}) tokenized_chat = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) return tokenized_chat # health_check.py import psutil from dataclasses import dataclass from typing import Dict, Any @dataclass class HealthStatus: status: str gpu_memory: Dict[str, float] cpu_usage: float ram_usage: float model_status: Dict[str, str] class HealthCheck: @staticmethod def check_gpu_memory() -> Dict[str, float]: if torch.cuda.is_available(): return { f"gpu_{i}": torch.cuda.memory_allocated(i) / 1024**3 for i in range(torch.cuda.device_count()) } return {} @staticmethod def check_system_resources() -> HealthStatus: return HealthStatus( status="healthy", gpu_memory=HealthCheck.check_gpu_memory(), cpu_usage=psutil.cpu_percent(), ram_usage=psutil.virtual_memory().percent, #TODO add more system resources like disk, network, etc. model_status={} # To be filled by the model manager ) # llama_generator.py from config.config import GenerationConfig, ModelConfig @observe() class LlamaGenerator(BaseGenerator): def __init__( self, llama_model_name: str, prm_model_path: str, device: Optional[str] = None, default_generation_config: Optional[GenerationConfig] = None, model_config: Optional[ModelConfig] = None, cache_size: int = 1000, max_batch_size: int = 32, # self.tokenizer = self.load_tokenizer(llama_model_name) # self.tokenizer = self.load_tokenizer(llama_model_name) # Add this line to initialize the tokenizer ): @observe() def load_model(self, model_name: str): # Code to load your model, e.g., Hugging Face's transformers library from transformers import AutoModelForCausalLM return AutoModelForCausalLM.from_pretrained(model_name) @observe() def load_tokenizer(self, model_name: str): # Load the tokenizer associated with the model from transformers import AutoTokenizer return AutoTokenizer.from_pretrained(model_name) self.tokenizer = load_tokenizer(llama_model_name) # Add this line to initialize the tokenizer super().__init__( llama_model_name, device, default_generation_config, model_config, cache_size, max_batch_size ) # Initialize models self.model_manager.load_model( "llama", llama_model_name, "llama", self.model_config ) self.model_manager.load_model( "prm", prm_model_path, "gguf", self.model_config ) self.prompt_builder = LlamaPromptTemplate() self._init_strategies() def _init_strategies(self): self.strategies = { "default": DefaultStrategy(), "majority_voting": MajorityVotingStrategy(), "best_of_n": BestOfN(), "beam_search": BeamSearch(), "dvts": DVT(), } def _get_generation_kwargs(self, config: GenerationConfig) -> Dict[str, Any]: """Get generation kwargs based on config.""" return { key: getattr(config, key) for key in [ "max_new_tokens", "temperature", "top_p", "top_k", "repetition_penalty", "length_penalty", "do_sample" ] if hasattr(config, key) } @observe() def generate_stream (self): return " NOt implememnted yet " @observe() def generate( self, prompt: str, model_kwargs: Dict[str, Any], strategy: str = "default", **kwargs ) -> str: """ Generate text based on a given strategy. Args: prompt (str): Input prompt for text generation. model_kwargs (Dict[str, Any]): Additional arguments for model generation. strategy (str): The generation strategy to use (default: "default"). **kwargs: Additional arguments passed to the strategy. Returns: str: Generated text response. Raises: ValueError: If the specified strategy is not available. """ # Validate that the strategy exists if strategy not in self.strategies: raise ValueError(f"Unknown strategy: {strategy}. Available strategies are: {list(self.strategies.keys())}") # Extract `generator` from kwargs if it exists to prevent duplication kwargs.pop("generator", None) # Call the selected strategy with the provided arguments return self.strategies[strategy].generate( generator=self, # The generator instance prompt=prompt, # The input prompt model_kwargs=model_kwargs, # Arguments for the model **kwargs # Any additional strategy-specific arguments ) @observe() def generate_with_context( self, context: str, user_input: str, chat_history: List[Tuple[str, str]], model_kwargs: Dict[str, Any], max_history_turns: int = 3, strategy: str = "default", num_samples: int = 5, depth: int = 3, breadth: int = 2, ) -> str: """Generate a response using context and chat history. Args: context (str): Context for the conversation user_input (str): Current user input chat_history (List[Tuple[str, str]]): List of (user, assistant) message pairs model_kwargs (dict): Additional arguments for model.generate() max_history_turns (int): Maximum number of history turns to include strategy (str): Generation strategy num_samples (int): Number of samples for applicable strategies depth (int): Depth for DVTS strategy breadth (int): Breadth for DVTS strategy Returns: str: Generated response """ prompt = self.prompt_builder.format( context, user_input, chat_history, max_history_turns ) return self.generate( generator=self, prompt=prompt, model_kwargs=model_kwargs, strategy=strategy, num_samples=num_samples, depth=depth, breadth=breadth ) def check_health(self) -> HealthStatus: """Check the health status of the generator.""" return self.health_check.check_system_resources() # TODO add model status ################### ################# from fastapi import FastAPI, HTTPException, BackgroundTasks, WebSocket, Depends from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field, ConfigDict from typing import List, Optional, Dict, Any, AsyncGenerator import asyncio import uuid from datetime import datetime import json from huggingface_hub import hf_hub_download from contextlib import asynccontextmanager class ChatMessage(BaseModel): """A single message in the chat history.""" role: str = Field( ..., description="Role of the message sender", examples=["user", "assistant"] ) content: str = Field(..., description="Content of the message") model_config = ConfigDict( json_schema_extra={ "example": { "role": "user", "content": "What is the capital of France?" } } ) class GenerationConfig(BaseModel): """Configuration for text generation.""" temperature: float = Field( 0.7, ge=0.0, le=2.0, description="Controls randomness in the output. Higher values (e.g., 0.8) make the output more random, lower values (e.g., 0.2) make it more focused and deterministic." ) max_new_tokens: int = Field( 100, ge=1, le=2048, description="Maximum number of tokens to generate" ) top_p: float = Field( 0.9, ge=0.0, le=1.0, description="Nucleus sampling parameter. Only tokens with cumulative probability < top_p are considered." ) top_k: int = Field( 50, ge=0, description="Only consider the top k tokens for text generation" ) strategy: str = Field( "default", description="Generation strategy to use", examples=["default", "majority_voting", "best_of_n", "beam_search", "dvts"] ) num_samples: int = Field( 5, ge=1, le=10, description="Number of samples to generate (used in majority_voting and best_of_n strategies)" ) class GenerationRequest(BaseModel): """Request model for text generation.""" context: Optional[str] = Field( None, description="Additional context to guide the generation", examples=["You are a helpful assistant skilled in Python programming"] ) messages: List[ChatMessage] = Field( ..., description="Chat history including the current message", min_items=1 ) config: Optional[GenerationConfig] = Field( None, description="Generation configuration parameters" ) stream: bool = Field( False, description="Whether to stream the response token by token" ) model_config = ConfigDict( json_schema_extra={ "example": { "context": "You are a helpful assistant", "messages": [ {"role": "user", "content": "What is the capital of France?"} ], "config": { "temperature": 0.7, "max_new_tokens": 100 }, "stream": False } } ) class GenerationResponse(BaseModel): """Response model for text generation.""" id: str = Field(..., description="Unique generation ID") content: str = Field(..., description="Generated text content") created_at: datetime = Field( default_factory=datetime.now, description="Timestamp of generation" ) # Model and cache management async def get_prm_model_path(): """Download and cache the PRM model.""" return await asyncio.to_thread( hf_hub_download, repo_id="tensorblock/Llama3.1-8B-PRM-Mistral-Data-GGUF", filename="Llama3.1-8B-PRM-Mistral-Data-Q4_K_M.gguf" ) # Initialize generator globally generator = None @asynccontextmanager async def lifespan(app: FastAPI): """Lifecycle management for the FastAPI application.""" # Startup: Initialize generator global generator try: prm_model_path = await get_prm_model_path() generator = LlamaGenerator( llama_model_name="meta-llama/Llama-3.2-1B-Instruct", prm_model_path=prm_model_path, default_generation_config=GenerationConfig( max_new_tokens=100, temperature=0.7 ) ) yield finally: # Shutdown: Clean up resources if generator: await asyncio.to_thread(generator.cleanup) # FastAPI application app = FastAPI( title="Inference Deluxe Service", description=""" A service for generating text using LLaMA models with various generation strategies. Generation Strategies: - default: Standard autoregressive generation - majority_voting: Generates multiple responses and selects the most common one - best_of_n: Generates multiple responses and selects the best based on a scoring metric - beam_search: Uses beam search for more coherent text generation - dvts: Dynamic vocabulary tree search for efficient generation """, version="1.0.0", lifespan=lifespan ) # CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) async def get_generator(): """Dependency to get the generator instance.""" if not generator: raise HTTPException( status_code=503, detail="Generator not initialized" ) return generator @app.post( "/generate", response_model=GenerationResponse, tags=["generation"], summary="Generate text response", response_description="Generated text with unique identifier" ) async def generate( request: GenerationRequest, generator: Any = Depends(get_generator) ): """ Generate a text response based on the provided context and chat history. """ try: chat_history = [(msg.role, msg.content) for msg in request.messages[:-1]] user_input = request.messages[-1].content # Extract or set defaults for additional arguments config = request.config or GenerationConfig() model_kwargs = { "temperature": config.temperature if hasattr(config, "temperature") else 0.7, "max_new_tokens": config.max_new_tokens if hasattr(config, "max_new_tokens") else 100, # Add other model kwargs as needed } # Explicitly pass additional required arguments response = await asyncio.to_thread( generator.generate_with_context, context=request.context or "", user_input=user_input, chat_history=chat_history, model_kwargs=model_kwargs, max_history_turns=config.max_history_turns if hasattr(config, "max_history_turns") else 3, strategy=config.strategy if hasattr(config, "strategy") else "default", num_samples=config.num_samples if hasattr(config, "num_samples") else 5, depth=config.depth if hasattr(config, "depth") else 3, breadth=config.breadth if hasattr(config, "breadth") else 2, ) return GenerationResponse( id=str(uuid.uuid4()), content=response ) except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.websocket("/generate/stream") async def generate_stream( websocket: WebSocket, generator: Any = Depends(get_generator) ): """ Stream generated text tokens over a WebSocket connection. The stream sends JSON messages with the following format: - During generation: {"token": "generated_token", "finished": false} - End of generation: {"token": "", "finished": true} - Error: {"error": "error_message"} """ await websocket.accept() try: while True: request_data = await websocket.receive_text() request = GenerationRequest.parse_raw(request_data) chat_history = [(msg.role, msg.content) for msg in request.messages[:-1]] user_input = request.messages[-1].content config = request.config or GenerationConfig() async for token in generator.generate_stream( prompt=generator.prompt_builder.format( context=request.context or "", user_input=user_input, chat_history=chat_history ), config=config ): await websocket.send_text(json.dumps({ "token": token, "finished": False })) await websocket.send_text(json.dumps({ "token": "", "finished": True })) except Exception as e: await websocket.send_text(json.dumps({ "error": str(e) })) finally: await websocket.close() if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)