"""HuggingFace service for interacting with the HuggingFace API. This service focuses on model discovery, search, and downloading from HuggingFace. """ import asyncio import os import shutil import tempfile from typing import Any, Dict, List, Optional, Tuple import aiohttp from starfish.common.logger import get_logger ##TODO we will need to move the dependencies of ollma to a seperate file so we can support other model hosting providers like vllm. but for now it is fine from starfish.llm.backend.ollama_adapter import delete_model as delete_ollama_model from starfish.llm.backend.ollama_adapter import is_model_available logger = get_logger(__name__) HF_API_BASE = "https://huggingface.co/api" ############################################# # HuggingFace Exception Types ############################################# class HuggingFaceError(Exception): """Base exception for HuggingFace-related errors.""" pass class HuggingFaceAuthError(HuggingFaceError): """Error raised when authentication is required but missing.""" pass class HuggingFaceModelNotFoundError(HuggingFaceError): """Error raised when a model is not found.""" pass class HuggingFaceAPIError(HuggingFaceError): """Error raised for general API errors.""" pass ############################################# # Core HuggingFace API Functions ############################################# def get_hf_token() -> Optional[str]: """Get HuggingFace API token from environment variable.""" return os.environ.get("HUGGING_FACE_HUB_TOKEN") async def _make_hf_request(url: str, params: Optional[Dict] = None, check_auth: bool = True) -> Tuple[bool, Any]: """Make a request to HuggingFace API with proper error handling. Args: url: API URL to request params: Optional query parameters check_auth: Whether to include auth token if available Returns: Tuple of (success, data/error_message) """ headers = {} if check_auth: token = get_hf_token() if token: headers["Authorization"] = f"Bearer {token}" try: async with aiohttp.ClientSession() as session: async with session.get(url, params=params, headers=headers) as response: if response.status == 200: return True, await response.json() elif response.status in (401, 403): return False, "Authentication required. Please set the {HUGGING_FACE_HUB_TOKEN} environment variable." else: return False, f"Request failed with status {response.status}" except Exception as e: return False, f"Request error: {str(e)}" async def list_hf_models(query: str = "", limit: int = 20) -> List[Dict[str, Any]]: """List/search models on HuggingFace. Args: query: Optional search query limit: Maximum number of results Returns: List of model info dictionaries """ params = {"limit": limit} if query: params["search"] = query success, result = await _make_hf_request(f"{HF_API_BASE}/models", params) if success: # Filter to only include models that are likely to work with Ollama # (supporting language models that might have GGUF variants) return [ { "id": model.get("id", ""), "name": model.get("modelId", model.get("id", "")), "downloads": model.get("downloads", 0), "likes": model.get("likes", 0), "tags": model.get("tags", []), "requires_auth": "gated" in model.get("tags", []) or model.get("private", False), } for model in result if "text-generation" in model.get("pipeline_tag", "") or any(tag in model.get("tags", []) for tag in ["llm", "gguf", "quantized"]) ] else: logger.error(f"Failed to list models: {result}") return [] async def get_imported_hf_models() -> List[str]: """Get list of HuggingFace models that have been imported to Ollama. Returns: List of model names in Ollama that originated from HuggingFace """ from starfish.llm.backend.ollama_adapter import list_models models = await list_models() return [model.get("name", "") for model in models if model.get("name", "").startswith("hf-")] async def check_model_exists(model_id: str) -> bool: """Check if a model exists on HuggingFace. Args: model_id: HuggingFace model ID (e.g., "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B") Returns: bool: True if model exists, False otherwise """ success, _ = await _make_hf_request(f"{HF_API_BASE}/models/{model_id}") return success async def find_gguf_files(model_id: str) -> List[Dict[str, Any]]: """Find GGUF files available for a HuggingFace model. Args: model_id: HuggingFace model ID Returns: List of file objects containing GGUF files """ # First try to check the main branch success, data = await _make_hf_request(f"{HF_API_BASE}/models/{model_id}/tree/main") if not success: # Try checking the master branch as fallback success, data = await _make_hf_request(f"{HF_API_BASE}/models/{model_id}/tree/master") if not success: return [] if not isinstance(data, list): logger.error(f"Unexpected data format from HuggingFace API: {type(data)}") return [] # First look for GGUF files gguf_files = [file for file in data if isinstance(file, dict) and file.get("path", "").lower().endswith(".gguf")] # If we found GGUF files, return them if gguf_files: logger.info(f"Found {len(gguf_files)} GGUF files for model {model_id}") return gguf_files # Otherwise, look for GGUF files in subdirectories for item in data: if isinstance(item, dict) and item.get("type") == "directory": dir_path = item.get("path") if dir_path: # Common directories where GGUF files are stored if any(keyword in dir_path.lower() for keyword in ["gguf", "quant", "quantized", "weights", "models"]): success, subdir_data = await _make_hf_request(f"{HF_API_BASE}/models/{model_id}/tree/main/{dir_path}") if success and isinstance(subdir_data, list): subdir_gguf_files = [file for file in subdir_data if isinstance(file, dict) and file.get("path", "").lower().endswith(".gguf")] if subdir_gguf_files: logger.info(f"Found {len(subdir_gguf_files)} GGUF files in subdirectory {dir_path}") gguf_files.extend(subdir_gguf_files) return gguf_files async def download_gguf_file(model_id: str, file_path: str, target_path: str) -> bool: """Download a GGUF file from HuggingFace. Args: model_id: HuggingFace model ID file_path: Path to the file within the repository target_path: Local path to save the file Returns: bool: True if download was successful, False otherwise """ url = f"https://huggingface.co/{model_id}/resolve/main/{file_path}" try: logger.info(f"Downloading {url} to {target_path}") # Ensure directory exists os.makedirs(os.path.dirname(target_path), exist_ok=True) headers = {} token = get_hf_token() if token: headers["Authorization"] = f"Bearer {token}" async with aiohttp.ClientSession() as session: async with session.get(url, headers=headers) as response: if response.status != 200: if response.status in (401, 403): return False # Auth error handled by caller logger.error(f"Error downloading GGUF file: {response.status}") return False # Download with progress reporting total_size = int(response.headers.get("Content-Length", 0)) chunk_size = 1024 * 1024 # 1MB chunks downloaded = 0 # Track last time progress was reported import time last_progress_time = 0 last_percentage = -1 # Always log the start logger.info("Download started: 0.0%") with open(target_path, "wb") as f: async for chunk in response.content.iter_chunked(chunk_size): f.write(chunk) downloaded += len(chunk) if total_size: current_time = time.time() progress = downloaded / total_size * 100 current_percentage = int(progress) # Log progress if 5+ seconds have passed or if we've hit a new percentage multiple of 25 if current_time - last_progress_time >= 5 or (current_percentage % 25 == 0 and current_percentage != last_percentage): logger.info(f"Download progress: {progress:.1f}%") last_progress_time = current_time last_percentage = current_percentage # Always log the completion logger.info(f"Download completed: {target_path}") return True except Exception as e: logger.error(f"Error downloading GGUF file: {e}") return False async def import_model_to_ollama(local_file_path: str, model_name: str) -> bool: """Import a GGUF file into Ollama. Args: local_file_path: Path to the downloaded GGUF file model_name: Name to give the model in Ollama Returns: bool: True if import was successful, False otherwise """ try: # Ensure Ollama bin exists ollama_bin = shutil.which("ollama") if not ollama_bin: logger.error("Ollama binary not found") return False # Create a temporary Modelfile with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: modelfile_path = f.name f.write(f"FROM {os.path.abspath(local_file_path)}\n") f.write('TEMPLATE "{.System}\\n\\n{.Prompt}"') logger.info(f"Created temporary Modelfile at {modelfile_path}") # Import the model using Ollama logger.info(f"Importing model into Ollama as {model_name}") process = await asyncio.create_subprocess_exec( ollama_bin, "create", model_name, "-f", modelfile_path, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) # Capture output stdout, stderr = await process.communicate() # Clean up temporary file os.unlink(modelfile_path) if process.returncode == 0: logger.info(f"Successfully imported model as {model_name}") return True else: logger.error(f"Failed to import model: {stderr.decode()}") return False except Exception as e: logger.error(f"Error importing model to Ollama: {e}") return False async def get_best_gguf_file(gguf_files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]: """Select the best GGUF file from a list of files. Prioritizes: 1. Smaller quantization (q4_K > q5_K > q8_0) 2. File size (prefers smaller files for same quantization level) 3. Avoid huge files unless necessary. Args: gguf_files: List of GGUF file objects Returns: The best GGUF file object or None if list is empty """ if not gguf_files: return None # Quantization priority (lower is better for typical use cases) # Prioritize middle-range quantization for balance of quality and size quant_priority = { "q4_k": 1, # Best balance for most use cases "q5_k": 2, # Good balance of quality and size "q4_0": 3, # Good for typical use "q4_1": 4, # Good for typical use "q3_k": 5, # More compression, less quality "q5_0": 6, "q5_1": 7, "q6_k": 8, "q2_k": 9, # Lowest quality but smallest "q8_0": 10, # High quality but larger "f16": 20, # Full precision, very large "f32": 30, # Full precision, extremely large } # Add a size penalty for very large files MAX_PREFERRED_SIZE = 4 * 1024 * 1024 * 1024 # 4GB # Extract quantization information for file in gguf_files: path = file.get("path", "").lower() size = file.get("size", 0) # Determine quantization level quant_score = 100 # Default high value for quant, priority in quant_priority.items(): if quant in path: quant_score = priority break # Store the score and size in the file object for sorting file["_quant_score"] = quant_score file["_size"] = size # First try to find models under the preferred size with good quantization preferred_files = [f for f in gguf_files if f.get("_size", 0) <= MAX_PREFERRED_SIZE] if preferred_files: # Sort by quantization score (lower is better) sorted_files = sorted(preferred_files, key=lambda x: x.get("_quant_score", 100)) else: # If all files are large, sort by quantization score sorted_files = sorted(gguf_files, key=lambda x: x.get("_quant_score", 100)) if sorted_files: selected = sorted_files[0] size_mb = selected.get("_size", 0) / (1024 * 1024) logger.info(f"Selected GGUF file: {selected.get('path')} ({size_mb:.1f} MB)") return selected return None async def download_best_gguf_for_model(model_id: str) -> Tuple[bool, str, Optional[str]]: """Download the best GGUF file for a HuggingFace model. This is a pure HuggingFace operation that doesn't directly depend on Ollama. Args: model_id: HuggingFace model ID Returns: Tuple of (success, message, local_file_path) Where local_file_path is the path to the downloaded file if successful, None otherwise Raises: HuggingFaceModelNotFoundError: If the model doesn't exist on HuggingFace HuggingFaceAuthError: If authentication is required but not provided HuggingFaceError: For other HuggingFace-related errors """ # Check if model exists if not await check_model_exists(model_id): error_msg = f"Model {model_id} not found on HuggingFace" logger.error(error_msg) raise HuggingFaceModelNotFoundError(error_msg) # Find GGUF files gguf_files = await find_gguf_files(model_id) if not gguf_files: error_msg = f"No GGUF files found for model {model_id}" logger.error(error_msg) raise HuggingFaceError(error_msg) # Select best GGUF file best_file = await get_best_gguf_file(gguf_files) if not best_file: error_msg = f"Could not select a suitable GGUF file for model {model_id}" logger.error(error_msg) raise HuggingFaceError(error_msg) # Create a temporary directory for download temp_dir = tempfile.mkdtemp() try: file_path = best_file.get("path") file_name = os.path.basename(file_path) local_file_path = os.path.join(temp_dir, file_name) # Download the GGUF file if not await download_gguf_file(model_id, file_path, local_file_path): # Check if it's an auth error token = get_hf_token() if not token: error_msg = f"Authentication required for model {model_id}. Please set the HUGGING_FACE_HUB_TOKEN environment variable." logger.error(error_msg) raise HuggingFaceAuthError(error_msg) else: error_msg = f"Failed to download GGUF file for model {model_id}" logger.error(error_msg) raise HuggingFaceError(error_msg) logger.info(f"Successfully downloaded GGUF file to {local_file_path}") return True, f"Successfully downloaded model {model_id}", local_file_path except Exception as e: # Clean up temp directory in case of error shutil.rmtree(temp_dir, ignore_errors=True) # Re-raise HuggingFace exceptions directly if isinstance(e, (HuggingFaceModelNotFoundError, HuggingFaceAuthError, HuggingFaceError)): raise # For other exceptions, wrap in HuggingFaceError logger.exception(f"Unexpected error downloading model {model_id}") raise HuggingFaceError(f"Unexpected error: {str(e)}") async def prepare_hf_model_for_ollama(model_id: str) -> Tuple[bool, str]: """Prepare a HuggingFace model for use with Ollama. This function bridges HuggingFace and Ollama services. Args: model_id: HuggingFace model ID Returns: Tuple of (success, message_or_ollama_name) Raises: HuggingFaceModelNotFoundError: If the model doesn't exist on HuggingFace HuggingFaceAuthError: If authentication is required but not provided HuggingFaceError: For other HuggingFace-related errors """ # Create a sanitized model name for Ollama ollama_name = f"hf-{model_id.replace('/', '-').lower()}" try: # Check if model is already imported in Ollama if await is_model_available(ollama_name): logger.info(f"Model {ollama_name} is already available in Ollama") return True, ollama_name # Download the best GGUF file success, message, local_file_path = await download_best_gguf_for_model(model_id) if not success or not local_file_path: # Re-raise the specific error based on the message if "not found" in message.lower(): raise HuggingFaceModelNotFoundError(message) elif "authentication" in message.lower(): raise HuggingFaceAuthError(message) else: raise HuggingFaceError(message) try: # Import to Ollama if not await import_model_to_ollama(local_file_path, ollama_name): raise HuggingFaceError(f"Failed to import model {model_id} to Ollama") logger.info(f"Successfully prepared HuggingFace model {model_id} as Ollama model {ollama_name}") return True, ollama_name finally: # Clean up the temp directory containing the downloaded file temp_dir = os.path.dirname(local_file_path) shutil.rmtree(temp_dir, ignore_errors=True) except (HuggingFaceModelNotFoundError, HuggingFaceAuthError, HuggingFaceError): # Let specific exceptions propagate up raise except Exception as e: logger.exception(f"Error preparing model {model_id} for Ollama") raise HuggingFaceError(f"Error preparing model: {str(e)}") async def ensure_hf_model_ready(model_id: str) -> Tuple[bool, str]: """Ensure a HuggingFace model is ready for use with Ollama. This function bridges HuggingFace and Ollama services. Args: model_id: HuggingFace model ID Returns: Tuple of (success, ollama_model_name) Raises: HuggingFaceModelNotFoundError: If the model doesn't exist on HuggingFace HuggingFaceAuthError: If authentication is required but not provided HuggingFaceError: For other HuggingFace-related errors """ from starfish.llm.backend.ollama_adapter import ( OllamaConnectionError, ensure_model_ready, ) # First ensure Ollama is running try: if not await ensure_model_ready(""): # Just make sure Ollama is running raise HuggingFaceError("Failed to start Ollama server") # Prepare the HuggingFace model for Ollama return await prepare_hf_model_for_ollama(model_id) except OllamaConnectionError as e: # Convert Ollama errors to HuggingFace errors for consistent error handling logger.error(f"Ollama connection error: {e}") raise HuggingFaceError(f"Ollama server error: {str(e)}") async def delete_hf_model(model_id: str) -> bool: """Delete a HuggingFace model from Ollama. This function bridges HuggingFace and Ollama services. Args: model_id: Either the HuggingFace model ID or the Ollama model name Returns: bool: True if deletion was successful """ # If this looks like a HF model ID, convert to Ollama name format if "/" in model_id and not model_id.startswith("hf-"): ollama_name = f"hf-{model_id.replace('/', '-').lower()}" else: ollama_name = model_id return await delete_ollama_model(ollama_name)