""" Models API Router - Explore and select models from HuggingFace Hub """ from fastapi import APIRouter, HTTPException, Query from pydantic import BaseModel, Field from typing import Optional, Dict, List, Any from datetime import datetime import logging import httpx from app.config import settings logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/models", tags=["Models"]) class ModelInfo(BaseModel): """Model information.""" id: str author: str model_name: str pipeline_tag: Optional[str] library_name: Optional[str] downloads: int = 0 likes: int = 0 private: bool = False tags: List[str] = [] created_at: Optional[datetime] last_modified: Optional[datetime] card_data: Optional[Dict[str, Any]] class ModelSearchResult(BaseModel): """Search result model.""" id: str author: str model_name: str pipeline_tag: Optional[str] downloads: int likes: int class ModelRecommendation(BaseModel): """Model recommendation for a task.""" model_id: str reason: str estimated_vram: str suggested_batch_size: int supports_peft: bool class ModelCompatibility(BaseModel): """Model compatibility check result.""" model_id: str compatible: bool issues: List[str] recommendations: List[str] max_sequence_length: int estimated_training_time: str requires_gpu: bool HF_API_URL = "https://huggingface.co/api" @router.get("/search", response_model=List[ModelSearchResult]) async def search_models( query: str = Query(..., min_length=1, description="Search query"), task: Optional[str] = Query(None, description="Filter by task/pipeline tag"), library: Optional[str] = Query(None, description="Filter by library"), author: Optional[str] = Query(None, description="Filter by author"), limit: int = Query(20, ge=1, le=100) ): """Search HuggingFace Hub for models.""" async with httpx.AsyncClient() as client: params = { "search": query, "limit": limit, "full": "false" } if task: params["pipeline_tag"] = task if library: params["library_name"] = library if author: params["author"] = author try: response = await client.get( f"{HF_API_URL}/models", params=params, timeout=30.0 ) response.raise_for_status() data = response.json() return [ ModelSearchResult( id=m.get("id", ""), author=m.get("author", ""), model_name=m.get("id", "").split("/")[-1] if "/" in m.get("id", "") else m.get("id", ""), pipeline_tag=m.get("pipeline_tag"), downloads=m.get("downloads", 0), likes=m.get("likes", 0) ) for m in data ] except Exception as e: logger.error(f"Search failed: {e}") raise HTTPException(status_code=500, detail=f"Search failed: {str(e)}") @router.get("/info/{model_id:path}", response_model=ModelInfo) async def get_model_info(model_id: str): """Get detailed information about a model.""" async with httpx.AsyncClient() as client: try: response = await client.get( f"{HF_API_URL}/models/{model_id}", timeout=30.0 ) response.raise_for_status() data = response.json() return ModelInfo( id=data.get("id", ""), author=data.get("author", ""), model_name=data.get("id", "").split("/")[-1] if "/" in data.get("id", "") else data.get("id", ""), pipeline_tag=data.get("pipeline_tag"), library_name=data.get("library_name"), downloads=data.get("downloads", 0), likes=data.get("likes", 0), private=data.get("private", False), tags=data.get("tags", []), created_at=datetime.fromisoformat(data["createdAt"].replace("Z", "+00:00")) if data.get("createdAt") else None, last_modified=datetime.fromisoformat(data["lastModified"].replace("Z", "+00:00")) if data.get("lastModified") else None, card_data=data.get("cardData") ) except httpx.HTTPStatusError as e: if e.response.status_code == 404: raise HTTPException(status_code=404, detail="Model not found") raise HTTPException(status_code=e.response.status_code, detail=f"API error: {e}") except Exception as e: logger.error(f"Failed to get model info: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.get("/recommend/{task_type}", response_model=List[ModelRecommendation]) async def get_model_recommendations(task_type: str): """Get model recommendations for a specific task.""" recommendations = { "causal-lm": [ ModelRecommendation(model_id="microsoft/phi-2", reason="Compact but powerful, excellent for fine-tuning", estimated_vram="8GB", suggested_batch_size=4, supports_peft=True), ModelRecommendation(model_id="google/gemma-2b", reason="Efficient 2B model, great for instruction tuning", estimated_vram="6GB", suggested_batch_size=8, supports_peft=True), ModelRecommendation(model_id="mistralai/Mistral-7B-v0.1", reason="Best-in-class 7B model", estimated_vram="16GB", suggested_batch_size=2, supports_peft=True), ModelRecommendation(model_id="meta-llama/Llama-3.2-3B", reason="Latest Llama architecture", estimated_vram="10GB", suggested_batch_size=4, supports_peft=True), ], "seq2seq": [ ModelRecommendation(model_id="google/flan-t5-base", reason="Versatile instruction-following model", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True), ModelRecommendation(model_id="google/flan-t5-large", reason="Larger variant for better quality", estimated_vram="8GB", suggested_batch_size=8, supports_peft=True), ModelRecommendation(model_id="facebook/bart-large", reason="Strong summarization model", estimated_vram="6GB", suggested_batch_size=8, supports_peft=True), ], "summarization": [ ModelRecommendation(model_id="facebook/bart-large-cnn", reason="Pre-trained for news summarization", estimated_vram="6GB", suggested_batch_size=8, supports_peft=True), ModelRecommendation(model_id="google/flan-t5-base", reason="Multi-purpose, handles various summarization tasks", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True), ModelRecommendation(model_id="mistralai/Mistral-7B-v0.1", reason="Large language model for advanced summarization", estimated_vram="16GB", suggested_batch_size=2, supports_peft=True), ], "token-classification": [ ModelRecommendation(model_id="dslim/bert-base-NER", reason="Strong NER model, 18 entity types", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True), ModelRecommendation(model_id="dbmdz/bert-large-cased-finetuned-conll03-english", reason="Large NER model for high accuracy", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True), ModelRecommendation(model_id="distilbert-base-cased", reason="Fast and efficient for fine-tuning", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True), ], "sequence-classification": [ ModelRecommendation(model_id="distilbert-base-uncased-finetuned-sst-2-english", reason="Sentiment analysis baseline", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True), ModelRecommendation(model_id="roberta-base", reason="Strong general-purpose classification model", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True), ModelRecommendation(model_id="microsoft/deberta-v3-base", reason="State-of-the-art classification accuracy", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True), ], "question-answering": [ ModelRecommendation(model_id="deepset/roberta-base-squad2", reason="SQuAD2.0 trained, handles unanswerable", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True), ModelRecommendation(model_id="distilbert-base-cased-distilled-squad", reason="Fast QA model", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True), ModelRecommendation(model_id="bert-large-uncased-whole-word-masking-finetuned-squad", reason="High accuracy QA", estimated_vram="4GB", suggested_batch_size=16, supports_peft=True), ], "translation": [ ModelRecommendation(model_id="Helsinki-NLP/opus-mt-en-de", reason="English to German", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True), ModelRecommendation(model_id="Helsinki-NLP/opus-mt-en-fr", reason="English to French", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True), ModelRecommendation(model_id="facebook/mbart-large-50-many-to-many-mmt", reason="Multilingual translation", estimated_vram="8GB", suggested_batch_size=8, supports_peft=True), ], "masked-lm": [ ModelRecommendation(model_id="bert-base-uncased", reason="Classic MLM model", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True), ModelRecommendation(model_id="roberta-base", reason="Modern MLM architecture", estimated_vram="2GB", suggested_batch_size=32, supports_peft=True), ModelRecommendation(model_id="distilbert-base-uncased", reason="Fast MLM training", estimated_vram="1GB", suggested_batch_size=64, supports_peft=True), ], } result = recommendations.get(task_type, []) if not result: result = recommendations.get("causal-lm", []) return result[:5] @router.get("/check/{model_id:path}", response_model=ModelCompatibility) async def check_model_compatibility( model_id: str, task: str = Query("causal-lm", description="Target task") ): """Check if a model is compatible for training with current setup.""" async with httpx.AsyncClient() as client: try: response = await client.get( f"{HF_API_URL}/models/{model_id}", timeout=30.0 ) response.raise_for_status() data = response.json() issues = [] recommendations = [] compatible = True library = data.get("library_name", "transformers") if library not in ["transformers", "peft", "adapter-transformers"]: issues.append(f"Unsupported library: {library}") compatible = False tags = data.get("tags", []) if "gated" in tags: recommendations.append("This is a gated model. Ensure you have access and set HF_TOKEN.") if data.get("private", False): recommendations.append("Private model requires authentication.") pipeline_tag = data.get("pipeline_tag", "") if "bert" in model_id.lower(): max_seq_len = 512 elif "llama" in model_id.lower(): max_seq_len = 4096 elif "mistral" in model_id.lower(): max_seq_len = 32768 elif "t5" in model_id.lower(): max_seq_len = 512 else: max_seq_len = 2048 task_pipeline_map = { "causal-lm": "text-generation", "seq2seq": "text2text-generation", "summarization": "summarization", "translation": "translation", "token-classification": "token-classification", "sequence-classification": "text-classification", "question-answering": "question-answering", "masked-lm": "fill-mask" } expected_pipeline = task_pipeline_map.get(task) if expected_pipeline and pipeline_tag and pipeline_tag != expected_pipeline: recommendations.append(f"Model pipeline tag ({pipeline_tag}) may not match task ({task}).") requires_gpu = False model_name_lower = model_id.lower() if any(x in model_name_lower for x in ["7b", "13b", "30b", "70b"]): issues.append("Large models (7B+) require GPU for practical training") requires_gpu = True compatible = False if "7b" in model_id.lower(): training_time = "Hours per epoch on GPU" elif "2b" in model_id.lower() or "3b" in model_id.lower(): training_time = "30-60 min per epoch on GPU" else: training_time = "10-30 min per epoch on GPU" return ModelCompatibility( model_id=model_id, compatible=compatible, issues=issues, recommendations=recommendations, max_sequence_length=max_seq_len, estimated_training_time=training_time, requires_gpu=requires_gpu ) except httpx.HTTPStatusError as e: if e.response.status_code == 404: raise HTTPException(status_code=404, detail="Model not found") raise HTTPException(status_code=e.response.status_code, detail=f"API error") except Exception as e: logger.error(f"Compatibility check failed: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.get("/popular") async def get_popular_models( task: Optional[str] = None, limit: int = Query(10, ge=1, le=50) ): """Get most popular models, optionally filtered by task.""" async with httpx.AsyncClient() as client: params = { "sort": "downloads", "direction": "-1", "limit": limit } if task: params["pipeline_tag"] = task try: response = await client.get( f"{HF_API_URL}/models", params=params, timeout=30.0 ) response.raise_for_status() return response.json() except Exception as e: logger.error(f"Failed to get popular models: {e}") raise HTTPException(status_code=500, detail=str(e)) @router.get("/tasks") async def get_supported_tasks(): """Get list of supported training tasks.""" return { "tasks": [ {"id": "causal-lm", "name": "Causal Language Modeling", "description": "Text generation, instruction tuning, chat models"}, {"id": "seq2seq", "name": "Sequence-to-Sequence", "description": "Input-output transformations"}, {"id": "summarization", "name": "Summarization", "description": "Text summarization models"}, {"id": "translation", "name": "Translation", "description": "Language translation models"}, {"id": "token-classification", "name": "Token Classification", "description": "NER, POS tagging, entity extraction"}, {"id": "sequence-classification", "name": "Sequence Classification", "description": "Sentiment, topic, intent classification"}, {"id": "question-answering", "name": "Question Answering", "description": "Extractive QA models"}, {"id": "masked-lm", "name": "Masked Language Modeling", "description": "BERT-style pretraining"}, ] }