| """ |
| 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"}, |
| ] |
| } |
