diff --git a/src/msquant/app/components/__init__.py b/src/msquant/app/components/__init__.py index 3278c09..d43c34c 100644 --- a/src/msquant/app/components/__init__.py +++ b/src/msquant/app/components/__init__.py @@ -1 +1,4 @@ """MSQuant application components.""" +from .hf_search_dialog import HFSearchDialog + +__all__ = ["HFSearchDialog"] diff --git a/src/msquant/app/components/hf_search_dialog.py b/src/msquant/app/components/hf_search_dialog.py new file mode 100644 index 0000000..ded5593 --- /dev/null +++ b/src/msquant/app/components/hf_search_dialog.py @@ -0,0 +1,301 @@ +"""HuggingFace search dialog component for models and datasets.""" +from typing import Callable, Optional, List +from nicegui import ui +from msquant.services import HuggingFaceService, HFSearchResult +import logging + +logger = logging.getLogger(__name__) + + +class HFSearchDialog: + """Dialog for searching HuggingFace models or datasets.""" + + def __init__( + self, + hf_service: HuggingFaceService, + search_type: str, + on_select: Callable[[str], None], + title: Optional[str] = None + ): + """ + Initialize the search dialog. + + Args: + hf_service: HuggingFace service instance + search_type: Type of search - 'model' or 'dataset' + on_select: Callback function when an item is selected, receives the item ID + title: Optional custom title for the dialog + """ + if search_type not in ['model', 'dataset']: + raise ValueError(f"search_type must be 'model' or 'dataset', got {search_type}") + + self.hf_service = hf_service + self.search_type = search_type + self.on_select = on_select + self.title = title or f"Search HuggingFace {search_type.capitalize()}s" + + self.dialog: ui.dialog + self.search_input: ui.input + self.results_container: ui.column + self.loading_spinner: ui.spinner + self.error_label: ui.label + self.sort_select: ui.select + self.direction_select: ui.select + self.search_results: List[HFSearchResult] = [] + + def show(self): + """Show the search dialog.""" + self.dialog = ui.dialog().props('maximized') + + with self.dialog, ui.card().classes('w-full h-full flex flex-col'): + # Header with title and close button + with ui.row().classes('w-full items-center justify-between mb-4'): + ui.label(self.title).classes('text-3xl font-bold') + ui.button(icon='close', on_click=self.dialog.close).props('flat round') + + # Search input and controls + with ui.row().classes('w-full gap-4 mb-4'): + self.search_input = ui.input( + placeholder=f'Search for {self.search_type}s...', + on_change=lambda: None # We'll use the search button + ).classes('flex-1').props('outlined clearable') + + # Bind Enter key to search + self.search_input.on('keydown.enter', self._perform_search) + + ui.button('Search', on_click=self._perform_search).props('color=primary') + + # Sort options + with ui.row().classes('w-full gap-4 mb-4'): + ui.label('Sort by:').classes('self-center') + self.sort_select = ui.select( + options=['downloads', 'likes', 'trending', 'modified', 'created'], + value='downloads', + label='Sort' + ).classes('w-48') + + self.direction_select = ui.select( + options=[ + {'label': 'Descending', 'value': -1}, + {'label': 'Ascending', 'value': 1} + ], + value=-1, + label='Direction' + ).classes('w-48').props('emit-value map-options') + + # Error message container + self.error_label = ui.label('').classes('text-red-500 mb-2') + self.error_label.set_visibility(False) + + # Loading spinner + with ui.row().classes('w-full justify-center'): + self.loading_spinner = ui.spinner(size='lg') + self.loading_spinner.set_visibility(False) + + # Results container with scrolling + with ui.scroll_area().classes('w-full flex-1'): + self.results_container = ui.column().classes('w-full gap-4') + + self.dialog.open() + + # Perform initial search with empty query to show popular items + self._perform_search() + + def _perform_search(self) -> None: + """Perform the search based on current input.""" + query = self.search_input.value.strip() if self.search_input.value else "" + sort = str(self.sort_select.value) if self.sort_select.value else 'downloads' + direction = int(self.direction_select.value) if self.direction_select.value is not None else -1 + + # Clear previous results and errors + self.results_container.clear() + self.error_label.set_text('') + self.error_label.set_visibility(False) + self.loading_spinner.set_visibility(True) + + try: + # Perform search + if self.search_type == 'model': + self.search_results = self.hf_service.search_models( + query=query, + limit=20, + sort=sort, + direction=direction + ) + else: # dataset + self.search_results = self.hf_service.search_datasets( + query=query, + limit=20, + sort=sort, + direction=direction + ) + + # Display results + with self.results_container: + if not self.search_results: + ui.label('No results found').classes('text-xl text-gray-500 text-center py-8') + else: + ui.label(f'Found {len(self.search_results)} results').classes('text-lg font-semibold mb-2') + + for result in self.search_results: + self._create_result_card(result) + + except Exception as e: + logger.error(f"Search failed: {e}") + self.error_label.set_text(f'Search failed: {str(e)}') + self.error_label.set_visibility(True) + + finally: + self.loading_spinner.set_visibility(False) + + def _create_result_card(self, result: HFSearchResult): + """Create a card for a single search result.""" + with ui.card().classes('w-full p-4 hover:shadow-lg cursor-pointer transition-shadow'): + with ui.row().classes('w-full items-start justify-between gap-4'): + # Left side - main info + with ui.column().classes('flex-1 gap-2'): + # Title with ID + with ui.row().classes('items-center gap-2'): + ui.label(result.id).classes('text-xl font-bold') + if result.author: + ui.label(f'by {result.author}').classes('text-gray-600') + + # Tags (show first 5) + if result.tags: + with ui.row().classes('gap-2 flex-wrap'): + for tag in result.tags[:5]: + ui.label(tag).classes('px-2 py-1 bg-blue-100 text-blue-800 rounded text-sm') + if len(result.tags) > 5: + ui.label(f'+{len(result.tags) - 5} more').classes('px-2 py-1 bg-gray-100 text-gray-600 rounded text-sm') + + # Additional info + with ui.row().classes('gap-4 text-sm text-gray-600'): + if result.downloads is not None: + ui.label(f'↓ {result.downloads:,} downloads') + if result.likes is not None: + ui.label(f'❤ {result.likes:,} likes') + if result.library_name: + ui.label(f'📚 {result.library_name}') + if result.pipeline_tag: + ui.label(f'🏷 {result.pipeline_tag}') + + # Right side - action buttons + with ui.column().classes('gap-2'): + ui.button( + 'Select', + on_click=lambda e=None, r=result: self._select_item(r) + ).props('color=primary') + + ui.button( + 'View Details', + on_click=lambda e=None, r=result: self._show_details(r) + ).props('outline') + + # Link to HuggingFace + hub_url = result.get_hub_url(self.search_type) + ui.link( + 'Open in HF', + hub_url, + new_tab=True + ).classes('text-blue-600 underline text-sm') + + def _select_item(self, result: HFSearchResult): + """Handle item selection.""" + try: + self.on_select(result.id) + self.dialog.close() + ui.notify(f'Selected: {result.id}', type='positive') + except Exception as e: + logger.error(f"Selection failed: {e}") + ui.notify(f'Failed to select item: {str(e)}', type='negative') + + def _show_details(self, result: HFSearchResult) -> None: + """Show detailed information about an item.""" + details_dialog = ui.dialog().props('maximized') + + with details_dialog, ui.card().classes('w-full h-full flex flex-col overflow-auto'): + # Header + with ui.row().classes('w-full items-center justify-between mb-4'): + ui.label(f'{self.search_type.capitalize()} Details').classes('text-3xl font-bold') + ui.button(icon='close', on_click=details_dialog.close).props('flat round') + + # Loading state + loading = ui.spinner(size='lg') + + try: + # Fetch detailed info + if self.search_type == 'model': + details = self.hf_service.get_model_details(result.id) + else: + details = self.hf_service.get_dataset_details(result.id) + + loading.set_visibility(False) + + # Display details + with ui.column().classes('w-full gap-4'): + # ID and basic info + ui.label(details['id']).classes('text-2xl font-bold') + + if details.get('author'): + with ui.row().classes('gap-2'): + ui.label('Author:').classes('font-semibold') + ui.label(details['author']) + + # Description + if details.get('description'): + ui.separator() + ui.label('Description').classes('text-xl font-bold') + ui.markdown(details['description']).classes('prose max-w-none') + else: + ui.label('No description available').classes('text-gray-500 italic') + + # Stats + ui.separator() + ui.label('Statistics').classes('text-xl font-bold') + with ui.grid(columns=2).classes('gap-4 w-full max-w-2xl'): + ui.label('Downloads:').classes('font-semibold') + ui.label(f"{details.get('downloads', 0):,}") + + ui.label('Likes:').classes('font-semibold') + ui.label(f"{details.get('likes', 0):,}") + + if details.get('library_name'): + ui.label('Library:').classes('font-semibold') + ui.label(details['library_name']) + + if details.get('pipeline_tag'): + ui.label('Pipeline:').classes('font-semibold') + ui.label(details['pipeline_tag']) + + ui.label('Last Modified:').classes('font-semibold') + ui.label(str(details.get('last_modified', 'N/A'))) + + # Tags + if details.get('tags'): + ui.separator() + ui.label('Tags').classes('text-xl font-bold') + with ui.row().classes('gap-2 flex-wrap'): + for tag in details['tags']: + ui.label(tag).classes('px-3 py-1 bg-blue-100 text-blue-800 rounded') + + # Actions + ui.separator() + with ui.row().classes('gap-4'): + ui.button( + 'Select This', + on_click=lambda: [self._select_item(result), details_dialog.close()] + ).props('color=primary size=lg') + + hub_url = result.get_hub_url(self.search_type) + ui.link( + 'Open in HuggingFace', + hub_url, + new_tab=True + ).classes('text-blue-600 underline text-lg') + + except Exception as e: + loading.set_visibility(False) + logger.error(f"Failed to load details: {e}") + ui.label(f'Failed to load details: {str(e)}').classes('text-red-500 text-xl') + + details_dialog.open() diff --git a/src/msquant/app/main.py b/src/msquant/app/main.py index 892f661..aca30d5 100644 --- a/src/msquant/app/main.py +++ b/src/msquant/app/main.py @@ -6,7 +6,7 @@ from msquant.app.pages import home, configure, monitor, results from msquant.app.components.layout import create_header -from msquant.services import JobService, StorageService +from msquant.services import JobService, StorageService, HuggingFaceService from msquant.core.monitoring import GPUMonitor @@ -18,11 +18,13 @@ hf_datasets_cache=os.environ.get("HF_DATASETS_CACHE", "/workspace/hf/datasets"), ) gpu_monitor = GPUMonitor(history_size=60) +hf_service = HuggingFaceService() # Store in app storage for access by pages app.storage.general['job_service'] = job_service app.storage.general['storage_service'] = storage_service app.storage.general['gpu_monitor'] = gpu_monitor +app.storage.general['hf_service'] = hf_service def init_routes(): @@ -36,7 +38,7 @@ def index_page(): @ui.page('/configure') def configure_page(): create_header() - configure.create_configure_page(job_service, storage_service) + configure.create_configure_page(job_service, storage_service, hf_service) @ui.page('/monitor') def monitor_page(): diff --git a/src/msquant/app/pages/configure.py b/src/msquant/app/pages/configure.py index 113e39b..f921912 100644 --- a/src/msquant/app/pages/configure.py +++ b/src/msquant/app/pages/configure.py @@ -1,10 +1,15 @@ """Configuration page for setting up quantization jobs.""" from nicegui import ui from msquant.core.quantizer import QuantizationConfig -from msquant.services import JobService, StorageService +from msquant.services import JobService, StorageService, HuggingFaceService +from msquant.app.components import HFSearchDialog -def create_configure_page(job_service: JobService, storage_service: StorageService): +def create_configure_page( + job_service: JobService, + storage_service: StorageService, + hf_service: HuggingFaceService +): """Create the configuration page.""" # Form state @@ -54,7 +59,32 @@ def start_quantization(): with ui.column().classes('w-full gap-4'): # Model settings ui.label('Model Settings').classes('text-2xl font-bold') - ui.input('Model ID', placeholder='e.g., meta-llama/Llama-3.1-8B').classes('w-full').bind_value(form_data, 'model_id') + + # Model ID input with search button + with ui.row().classes('w-full gap-2'): + model_input = ui.input( + 'Model ID', + placeholder='e.g., meta-llama/Llama-3.1-8B' + ).classes('flex-1').bind_value(form_data, 'model_id') + + def show_model_search(): + def on_model_select(model_id: str): + form_data['model_id'] = model_id + model_input.update() + + search_dialog = HFSearchDialog( + hf_service=hf_service, + search_type='model', + on_select=on_model_select, + title='Search HuggingFace Models' + ) + search_dialog.show() + + ui.button( + 'Search Models', + on_click=show_model_search, + icon='search' + ).props('color=secondary') ui.separator() @@ -67,7 +97,32 @@ def start_quantization(): # Calibration settings ui.label('Calibration Settings').classes('text-2xl font-bold') - ui.input('Calibration Dataset', placeholder='e.g., wikitext').classes('w-full').bind_value(form_data, 'calib_dataset') + + # Calibration Dataset input with search button + with ui.row().classes('w-full gap-2'): + dataset_input = ui.input( + 'Calibration Dataset', + placeholder='e.g., wikitext' + ).classes('flex-1').bind_value(form_data, 'calib_dataset') + + def show_dataset_search(): + def on_dataset_select(dataset_id: str): + form_data['calib_dataset'] = dataset_id + dataset_input.update() + + search_dialog = HFSearchDialog( + hf_service=hf_service, + search_type='dataset', + on_select=on_dataset_select, + title='Search HuggingFace Datasets' + ) + search_dialog.show() + + ui.button( + 'Search Datasets', + on_click=show_dataset_search, + icon='search' + ).props('color=secondary') ui.input('Dataset Config', placeholder='e.g., wikitext-2-raw-v1 (optional)').classes('w-full').bind_value(form_data, 'calib_config') ui.select( ['train', 'test', 'validation', 'train[:512]', 'test[:512]'], diff --git a/src/msquant/services/__init__.py b/src/msquant/services/__init__.py index 4b380d1..4a153eb 100644 --- a/src/msquant/services/__init__.py +++ b/src/msquant/services/__init__.py @@ -1,5 +1,6 @@ """MSQuant services.""" from .jobs import JobService from .storage import StorageService +from .huggingface import HuggingFaceService, HFSearchResult -__all__ = ["JobService", "StorageService"] +__all__ = ["JobService", "StorageService", "HuggingFaceService", "HFSearchResult"] diff --git a/src/msquant/services/huggingface.py b/src/msquant/services/huggingface.py new file mode 100644 index 0000000..f37cb34 --- /dev/null +++ b/src/msquant/services/huggingface.py @@ -0,0 +1,297 @@ +"""HuggingFace Hub API service for searching models and datasets.""" +from typing import List, Dict, Any, Optional, Literal +from dataclasses import dataclass +from huggingface_hub import list_models, list_datasets, model_info, dataset_info +from huggingface_hub.hf_api import ModelInfo, DatasetInfo +import logging + +logger = logging.getLogger(__name__) + + +@dataclass +class HFSearchResult: + """Represents a search result from HuggingFace Hub.""" + id: str + author: Optional[str] + description: Optional[str] + downloads: Optional[int] + likes: Optional[int] + tags: List[str] + last_modified: Optional[str] + created_at: Optional[str] + library_name: Optional[str] = None # For models + pipeline_tag: Optional[str] = None # For models + + def get_hub_url(self, repo_type: str) -> str: + """Get the HuggingFace Hub URL for this item.""" + return f"https://huggingface.co/{repo_type}s/{self.id}" + + +class HuggingFaceService: + """Service for interacting with HuggingFace Hub API.""" + + def __init__(self): + """Initialize the HuggingFace service.""" + self._cache: Dict[str, Any] = {} + + def search_models( + self, + query: str = "", + limit: int = 20, + sort: str = "downloads", + direction: int = -1 + ) -> List[HFSearchResult]: + """ + Search for models on HuggingFace Hub. + + Args: + query: Search query string + limit: Maximum number of results to return (max 100) + sort: Property to sort by (downloads, likes, trending, etc.) + direction: Sort direction (-1 for descending, 1 for ascending) + + Returns: + List of HFSearchResult objects + + Raises: + ValueError: If parameters are invalid + RuntimeError: If API call fails + """ + # Validate inputs + if not isinstance(query, str): + raise ValueError(f"query must be a string, got {type(query)}") + if not isinstance(limit, int) or limit < 1 or limit > 100: + raise ValueError(f"limit must be an integer between 1 and 100, got {limit}") + if sort not in ["downloads", "likes", "trending", "author", "id", "created", "modified"]: + raise ValueError(f"Invalid sort parameter: {sort}") + if direction not in [-1, 1]: + raise ValueError(f"direction must be -1 or 1, got {direction}") + + try: + logger.info(f"Searching models: query={query}, limit={limit}, sort={sort}") + + # Convert direction: -1 stays as -1, 1 becomes None (for ascending) + hf_direction: Optional[Literal[-1]] = -1 if direction == -1 else None + + # Use huggingface_hub library to search models + models = list_models( + search=query if query else None, + limit=limit, + sort=sort, + direction=hf_direction + ) + + results = [] + for model in models: + try: + result = self._model_to_search_result(model) + results.append(result) + except Exception as e: + logger.warning(f"Failed to process model {getattr(model, 'id', 'unknown')}: {e}") + continue + + logger.info(f"Found {len(results)} models") + return results + + except Exception as e: + logger.error(f"Failed to search models: {e}") + raise RuntimeError(f"Failed to search models: {str(e)}") from e + + def search_datasets( + self, + query: str = "", + limit: int = 20, + sort: str = "downloads", + direction: int = -1 + ) -> List[HFSearchResult]: + """ + Search for datasets on HuggingFace Hub. + + Args: + query: Search query string + limit: Maximum number of results to return (max 100) + sort: Property to sort by (downloads, likes, trending, etc.) + direction: Sort direction (-1 for descending, 1 for ascending) + + Returns: + List of HFSearchResult objects + + Raises: + ValueError: If parameters are invalid + RuntimeError: If API call fails + """ + # Validate inputs + if not isinstance(query, str): + raise ValueError(f"query must be a string, got {type(query)}") + if not isinstance(limit, int) or limit < 1 or limit > 100: + raise ValueError(f"limit must be an integer between 1 and 100, got {limit}") + if sort not in ["downloads", "likes", "trending", "author", "id", "created", "modified"]: + raise ValueError(f"Invalid sort parameter: {sort}") + if direction not in [-1, 1]: + raise ValueError(f"direction must be -1 or 1, got {direction}") + + try: + logger.info(f"Searching datasets: query={query}, limit={limit}, sort={sort}") + + # Convert direction: -1 stays as -1, 1 becomes None (for ascending) + hf_direction: Optional[Literal[-1]] = -1 if direction == -1 else None + + # Use huggingface_hub library to search datasets + datasets = list_datasets( + search=query if query else None, + limit=limit, + sort=sort, + direction=hf_direction + ) + + results = [] + for dataset in datasets: + try: + result = self._dataset_to_search_result(dataset) + results.append(result) + except Exception as e: + logger.warning(f"Failed to process dataset {getattr(dataset, 'id', 'unknown')}: {e}") + continue + + logger.info(f"Found {len(results)} datasets") + return results + + except Exception as e: + logger.error(f"Failed to search datasets: {e}") + raise RuntimeError(f"Failed to search datasets: {str(e)}") from e + + def get_model_details(self, model_id: str) -> Dict[str, Any]: + """ + Get detailed information about a specific model. + + Args: + model_id: The model repository ID (e.g., 'meta-llama/Llama-3.1-8B') + + Returns: + Dictionary containing detailed model information + + Raises: + ValueError: If model_id is invalid + RuntimeError: If API call fails + """ + if not isinstance(model_id, str) or not model_id.strip(): + raise ValueError("model_id must be a non-empty string") + + try: + logger.info(f"Fetching model details for: {model_id}") + info = model_info(model_id, files_metadata=False) + + # Extract card data for description + card_data = getattr(info, 'card_data', None) or getattr(info, 'cardData', None) + description = "" + + if card_data: + # Try to get description from card_data + if hasattr(card_data, 'get'): + description = card_data.get('description', '') or card_data.get('model-index', {}).get('results', [{}])[0].get('task', {}).get('name', '') + elif hasattr(card_data, 'to_dict'): + card_dict = card_data.to_dict() + description = card_dict.get('description', '') + + # Fallback to checking for description attribute + if not description: + description = getattr(info, 'description', '') + + return { + 'id': info.id, + 'author': getattr(info, 'author', None), + 'description': description, + 'downloads': getattr(info, 'downloads', 0), + 'likes': getattr(info, 'likes', 0), + 'tags': getattr(info, 'tags', []), + 'library_name': getattr(info, 'library_name', None), + 'pipeline_tag': getattr(info, 'pipeline_tag', None), + 'last_modified': str(getattr(info, 'lastModified', '')), + 'created_at': str(getattr(info, 'created_at', '')), + 'card_data': card_data, + } + + except Exception as e: + logger.error(f"Failed to get model details for {model_id}: {e}") + raise RuntimeError(f"Failed to get model details: {str(e)}") from e + + def get_dataset_details(self, dataset_id: str) -> Dict[str, Any]: + """ + Get detailed information about a specific dataset. + + Args: + dataset_id: The dataset repository ID (e.g., 'wikitext') + + Returns: + Dictionary containing detailed dataset information + + Raises: + ValueError: If dataset_id is invalid + RuntimeError: If API call fails + """ + if not isinstance(dataset_id, str) or not dataset_id.strip(): + raise ValueError("dataset_id must be a non-empty string") + + try: + logger.info(f"Fetching dataset details for: {dataset_id}") + info = dataset_info(dataset_id, files_metadata=False) + + # Extract card data for description + card_data = getattr(info, 'card_data', None) or getattr(info, 'cardData', None) + description = "" + + if card_data: + # Try to get description from card_data + if hasattr(card_data, 'get'): + description = card_data.get('description', '') or card_data.get('dataset_summary', '') + elif hasattr(card_data, 'to_dict'): + card_dict = card_data.to_dict() + description = card_dict.get('description', '') or card_dict.get('dataset_summary', '') + + # Fallback to checking for description attribute + if not description: + description = getattr(info, 'description', '') + + return { + 'id': info.id, + 'author': getattr(info, 'author', None), + 'description': description, + 'downloads': getattr(info, 'downloads', 0), + 'likes': getattr(info, 'likes', 0), + 'tags': getattr(info, 'tags', []), + 'last_modified': str(getattr(info, 'lastModified', '')), + 'created_at': str(getattr(info, 'created_at', '')), + 'card_data': card_data, + } + + except Exception as e: + logger.error(f"Failed to get dataset details for {dataset_id}: {e}") + raise RuntimeError(f"Failed to get dataset details: {str(e)}") from e + + def _model_to_search_result(self, model: ModelInfo) -> HFSearchResult: + """Convert ModelInfo to HFSearchResult.""" + return HFSearchResult( + id=model.id, + author=getattr(model, 'author', None), + description=None, # Description not available in list results + downloads=getattr(model, 'downloads', 0), + likes=getattr(model, 'likes', 0), + tags=getattr(model, 'tags', []) or [], + last_modified=str(getattr(model, 'lastModified', '')), + created_at=str(getattr(model, 'created_at', '')), + library_name=getattr(model, 'library_name', None), + pipeline_tag=getattr(model, 'pipeline_tag', None), + ) + + def _dataset_to_search_result(self, dataset: DatasetInfo) -> HFSearchResult: + """Convert DatasetInfo to HFSearchResult.""" + return HFSearchResult( + id=dataset.id, + author=getattr(dataset, 'author', None), + description=getattr(dataset, 'description', None), + downloads=getattr(dataset, 'downloads', 0), + likes=getattr(dataset, 'likes', 0), + tags=getattr(dataset, 'tags', []) or [], + last_modified=str(getattr(dataset, 'lastModified', '')), + created_at=str(getattr(dataset, 'created_at', '')), + )