""" Research Agent - Searches HuggingFace for relevant models """ import asyncio import aiohttp import re from datetime import datetime from typing import List, Optional from huggingface_hub import HfApi from src.models.schemas import UserRequirements, ModelMetadata, TaskType class ResearchAgent: """ Searches HuggingFace Hub for models matching user requirements. This agent queries the HuggingFace API, fetches model details, and creates ModelMetadata objects for candidate models. """ # Pipeline tags for different tasks PIPELINE_TAGS = { # Translation TaskType.TRANSLATION: "translation", # Speech TaskType.TEXT_TO_SPEECH: "text-to-speech", TaskType.SPEECH_TO_TEXT: "automatic-speech-recognition", # LLMs TaskType.TEXT_GENERATION: "text-generation", TaskType.CHAT: "text-generation", TaskType.INSTRUCTION_FOLLOWING: "text-generation", TaskType.CODE_GENERATION: "text-generation", TaskType.QUESTION_ANSWERING: "question-answering", TaskType.SUMMARIZATION: "summarization", # OCR TaskType.OCR: "image-to-text", TaskType.DOCUMENT_UNDERSTANDING: "document-question-answering", #Legacy TaskType.TEXT_CLASSIFICATION: "text-classification", TaskType.NAMED_ENTITY_RECOGNITION: "token-classification", TaskType.IMAGE_CLASSIFICATION: "image-classification", TaskType.OBJECT_DETECTION: "object-detection", TaskType.ZERO_SHOT_CLASSIFICATION: "zero-shot-classification" } def __init__(self): self.api = HfApi() async def search_models(self, requirements: UserRequirements, top_k: int = 20) -> List[ModelMetadata]: """ Search for models matching the requirements. Args: requirements: UserRequirements object top_k: Maximum number of models to return Returns: List of ModelMetadata objects """ task = requirements.task_type pipeline_tag = self.PIPELINE_TAGS.get(task) print(f" Searching for {task.value} models (pipeline: {pipeline_tag})...") try: models = [] # Method 1: Search by pipeline_tag filter (current API) if pipeline_tag: try: print(f" Method 1: Using pipeline_tag filter...") models = list(self.api.list_models( filter=f"pipeline_tag:{pipeline_tag}", sort="downloads", limit=top_k * 2 )) print(f" Found {len(models)} models") except Exception as e: print(f" Method 1 failed: {e}") models = [] # Method 1.5: For translation, try specific translation tags if not models and task == TaskType.TRANSLATION: try: print(f" Method 1.5: Trying specific translation tags...") # Try common translation model patterns search_terms = ["translation", "mbart", "nllb", "m2m", "opus"] for term in search_terms: try: term_models = list(self.api.list_models( search=term, sort="downloads", limit=top_k )) models.extend(term_models) print(f" Found {len(term_models)} models with '{term}'") except: continue # Remove duplicates unique_ids = set() unique_models = [] for m in models: if m.modelId not in unique_ids: unique_ids.add(m.modelId) unique_models.append(m) models = unique_models print(f" Total unique models after search: {len(models)}") except Exception as e: print(f" Method 1.5 failed: {e}") # Method 2: Try without filter (get popular models) if not models: try: print(f" Method 2: Getting popular models...") models = list(self.api.list_models( sort="downloads", limit=top_k * 2 )) print(f" Found {len(models)} models") except Exception as e: print(f" Method 2 failed: {e}") models = [] # Method 3: Try with search parameter if not models and pipeline_tag: try: print(f" Method 3: Using search parameter...") models = list(self.api.list_models( search=pipeline_tag, sort="downloads", limit=top_k * 2 )) print(f" Found {len(models)} models") except Exception as e: print(f" Method 3 failed: {e}") models = [] print(f" Total candidate models: {len(models)}") if not models: print(" No models found from HuggingFace API") return [] # Fetch detailed metadata model_details = [] for i, model in enumerate(models[:top_k]): try: print(f" Processing {i+1}/{min(len(models), top_k)}: {model.modelId}") metadata = await self._fetch_model_details(model, requirements) if metadata: model_details.append(metadata) print(f" Added to candidates") await asyncio.sleep(0.1) except Exception as e: print(f" Error: {e}") continue print(f" Found {len(model_details)} suitable models") return model_details except Exception as e: print(f" Error searching models: {e}") import traceback traceback.print_exc() return [] async def _fetch_model_details(self, model, requirements: UserRequirements) -> Optional[ModelMetadata]: """Fetch detailed information for a specific model""" try: # Get model info from HuggingFace model_info = self.api.model_info(model.modelId) # Fetch model card content (README.md) model_card = await self._fetch_model_card(model.modelId) # Extract performance metrics from model card performance_metrics = self._extract_performance_metrics(model_card) # Estimate model size model_size = self._estimate_model_size(model_info) # Filter by size constraint if specified if requirements.max_model_size_gb and model_size: if model_size > requirements.max_model_size_gb: return None # Extract language tags languages = ["en"] # Default if hasattr(model_info, 'tags'): for tag in model_info.tags: if tag.startswith("language:"): languages = [tag.replace("language:", "")] break # Get license from card data license_info = "unknown" if hasattr(model_info, 'cardData') and model_info.cardData: license_info = model_info.cardData.get("license", "unknown") # Get tags tags = getattr(model_info, 'tags', []) # Get pipeline tag pipeline_tag = getattr(model_info, 'pipeline_tag', None) # Get LLM-specific info llm_info = self._extract_llm_info(model_info, model_card) # Get OCR-specific info ocr_info = self._extract_ocr_info(model_info, model_card) # Determine task type from pipeline tag task_type = self._map_pipeline_to_task(pipeline_tag) or requirements.task_type # Extract source and target languages for translation source_languages = self._extract_source_languages(model_info, model_card) target_languages = self._extract_target_languages(model_info, model_card) # Extract voice count for TTS voice_count = self._extract_voice_count(model_card) # Extract sample rate for audio models sample_rate = self._extract_sample_rate(model_card) # Extract WER for STT wer_score = self._extract_wer(model_card) return ModelMetadata( model_id=model.modelId, task_type=task_type, downloads=getattr(model_info, 'downloads', 0) or 0, likes=getattr(model_info, 'likes', 0) or 0, last_modified=getattr(model_info, 'lastModified', datetime.now()), license=license_info, model_size=model_size, languages=languages, tags=tags, pipeline_tag=pipeline_tag, base_model=getattr(model_info, 'base_model', None), finetuned_from=getattr(model_info, 'finetuned_from', None), performance_metrics=performance_metrics, hardware_requirements=self._extract_hardware_info(model_card), model_card_content=model_card, # Additional fields source_languages=source_languages, target_languages=target_languages, voice_count=voice_count, sample_rate=sample_rate, wer_score=wer_score, context_length=llm_info.get("context_length"), has_chat_template=llm_info.get("has_chat_template", False), supports_function_calling=llm_info.get("supports_function_calling", False), supports_code=llm_info.get("supports_code", False), supports_instruction=llm_info.get("supports_instruction", False), quantization_supported=llm_info.get("quantization_supported", []), supports_handwriting=ocr_info.get("supports_handwriting", False), supports_layout=ocr_info.get("supports_layout", False), supports_tables=ocr_info.get("supports_tables", False), supports_formulas=ocr_info.get("supports_formulas", False), supported_image_formats=ocr_info.get("supported_image_formats", ["jpg", "png"]) ) except Exception as e: print(f" Error fetching details: {e}") return None async def _fetch_model_card(self, model_id: str) -> str: """Fetch model card content from HuggingFace""" try: card_url = f"https://huggingface.co/{model_id}/raw/main/README.md" async with aiohttp.ClientSession() as session: async with session.get(card_url, timeout=10) as response: if response.status == 200: return await response.text() return "" except Exception: return "" def _extract_performance_metrics(self, model_card: str) -> dict: """Extract performance metrics from model card text""" metrics = {} # Common metrics to look for metric_patterns = { "accuracy": r"accuracy[\s:]*([\d\.]+)%?", "f1": r"f1[\s:]*([\d\.]+)", "bleu": r"bleu[\s:]*([\d\.]+)", "rouge": r"rouge[\s:]*([\d\.]+)", "perplexity": r"perplexity[\s:]*([\d\.]+)", "precision": r"precision[\s:]*([\d\.]+)", "recall": r"recall[\s:]*([\d\.]+)" } for metric, pattern in metric_patterns.items(): matches = re.findall(pattern, model_card.lower()) if matches: try: metrics[metric] = float(matches[0]) except ValueError: pass return metrics def _estimate_model_size(self, model_info) -> Optional[float]: """Estimate model size in GB""" try: # Try to get from config if hasattr(model_info, 'config') and model_info.config: param_size = model_info.config.get("num_parameters", 0) if param_size: # Rough estimate: 4 bytes per parameter (float32) size_gb = (param_size * 4) / (1024 ** 3) return size_gb # Alternative: look for safetensors files if hasattr(model_info, 'siblings'): total_size = 0 for sibling in model_info.siblings: if hasattr(sibling, 'rfilename') and sibling.rfilename.endswith(('.safetensors', '.bin')): if hasattr(sibling, 'size'): total_size += sibling.size if total_size > 0: return total_size / (1024 ** 3) # Convert to GB except Exception: pass # Default size based on model name patterns model_id = model_info.modelId.lower() if hasattr(model_info, 'modelId') else "" if any(x in model_id for x in ['tiny', 'mini', 'albert']): return 0.05 # 50MB elif any(x in model_id for x in ['small', 'distilbert']): return 0.2 # 200MB elif any(x in model_id for x in ['base', 'bert-base']): return 0.5 # 500MB elif any(x in model_id for x in ['large', 'bert-large']): return 1.5 # 1.5GB elif any(x in model_id for x in ['xl', 'gpt2-xl']): return 3.0 # 3GB return None def _extract_hardware_info(self, model_card: str) -> dict: """Extract hardware requirements from model card""" hardware = { "cpu_compatible": True, # Assume CPU compatible by default "gpu_required": False, "tpu_compatible": False, "min_ram_gb": 4 # Default assumption } card_lower = model_card.lower() if "gpu" in card_lower and "no gpu" not in card_lower: hardware["gpu_required"] = True if "tpu" in card_lower: hardware["tpu_compatible"] = True # Look for RAM requirements ram_match = re.search(r'(\d+)\s*gb?\s*ram', card_lower) if ram_match: hardware["min_ram_gb"] = int(ram_match.group(1)) return hardware def _extract_source_languages(self, model_info, model_card: str) -> List[str]: """Extract source languages for translation models""" languages = [] if hasattr(model_info, 'cardData') and model_info.cardData: src_langs = model_info.cardData.get("src_lang", []) if src_langs: if isinstance(src_langs, str): languages = [src_langs] elif isinstance(src_langs, list): languages = src_langs return languages or ["en"] def _extract_target_languages(self, model_info, model_card: str) -> List[str]: """Extract target languages for translation models""" languages = [] if hasattr(model_info, 'cardData') and model_info.cardData: tgt_langs = model_info.cardData.get("tgt_lang", []) if tgt_langs: if isinstance(tgt_langs, str): languages = [tgt_langs] elif isinstance(tgt_langs, list): languages = tgt_langs return languages or ["en"] def _extract_voice_count(self, model_card: str) -> int: """Extract number of voices for TTS models""" patterns = [ r'(\d+)\s+voices?', r'voices?:?\s*(\d+)', r'multi-voice.*?(\d+)', ] for pattern in patterns: match = re.search(pattern, model_card.lower()) if match: try: return int(match.group(1)) except: pass return 0 def _extract_sample_rate(self, model_card: str) -> Optional[int]: """Extract sample rate for audio models""" patterns = [ r'(\d+)\s*[kK]?[hH][zZ]', r'sample rate:?\s*(\d+)', ] for pattern in patterns: match = re.search(pattern, model_card.lower()) if match: try: rate = int(match.group(1)) if 'k' in match.group(0).lower(): rate *= 1000 return rate except: pass return None def _extract_wer(self, model_card: str) -> Optional[float]: """Extract Word Error Rate for STT models""" patterns = [ r'wer:?\s*([\d.]+)%?', r'word error rate:?\s*([\d.]+)%?', r'wer[\s=]+([\d.]+)', ] for pattern in patterns: match = re.search(pattern, model_card.lower()) if match: try: return float(match.group(1)) except: pass return None def _extract_llm_info(self, model_info, model_card: str) -> dict: """Extract LLM-specific information""" info = { "context_length": 2048, # default "has_chat_template": False, "supports_function_calling": False, "supports_code": False, "supports_instruction": False, "quantization_supported": [] } card_lower = model_card.lower() # Check context length context_patterns = [ r'context length:?\s*(\d+)[kK]?', r'max[ _]?length:?\s*(\d+)[kK]?', r'(\d+)[kK]\s*(?:context|tokens)' ] for pattern in context_patterns: match = re.search(pattern, card_lower) if match: val = int(match.group(1)) if 'k' in match.group(0).lower(): info["context_length"] = val * 1024 else: info["context_length"] = val break # Check for chat template info["has_chat_template"] = "chat template" in card_lower or "conversation" in card_lower # Check for function calling info["supports_function_calling"] = any(x in card_lower for x in ["function calling", "tools", "function call", "tool use"]) # Check for code generation info["supports_code"] = any(x in card_lower for x in ["code generation", "programming", "python", "javascript"]) # Check for instruction following info["supports_instruction"] = "instruction" in card_lower # Check quantization support if "4bit" in card_lower or "4-bit" in card_lower: info["quantization_supported"].append("4bit") if "8bit" in card_lower or "8-bit" in card_lower: info["quantization_supported"].append("8bit") return info def _extract_ocr_info(self, model_info, model_card: str) -> dict: """Extract OCR-specific information""" info = { "supports_handwriting": False, "supports_layout": False, "supports_tables": False, "supports_formulas": False, "supported_image_formats": ["jpg", "png"] # default } card_lower = model_card.lower() # Check for handwriting info["supports_handwriting"] = "handwriting" in card_lower or "handwritten" in card_lower # Check for layout analysis info["supports_layout"] = any(x in card_lower for x in ["layout", "paragraph", "document structure"]) # Check for table extraction info["supports_tables"] = any(x in card_lower for x in ["table", "spreadsheet", "tabular"]) # Check for formula recognition info["supports_formulas"] = any(x in card_lower for x in ["formula", "equation", "math", "latex"]) # Check image formats formats = [] for fmt in ["jpg", "jpeg", "png", "tiff", "bmp", "pdf"]: if fmt in card_lower: formats.append(fmt) if formats: info["supported_image_formats"] = formats return info def _map_pipeline_to_task(self, pipeline_tag: Optional[str]) -> Optional[TaskType]: """Map HuggingFace pipeline tag to our TaskType""" mapping = { "translation": TaskType.TRANSLATION, "text-to-speech": TaskType.TEXT_TO_SPEECH, "automatic-speech-recognition": TaskType.SPEECH_TO_TEXT, "text-generation": TaskType.TEXT_GENERATION, "question-answering": TaskType.QUESTION_ANSWERING, "summarization": TaskType.SUMMARIZATION, "image-to-text": TaskType.OCR, "document-question-answering": TaskType.DOCUMENT_UNDERSTANDING, "text-classification": TaskType.TEXT_CLASSIFICATION, "token-classification": TaskType.NAMED_ENTITY_RECOGNITION, "image-classification": TaskType.IMAGE_CLASSIFICATION, "object-detection": TaskType.OBJECT_DETECTION, "zero-shot-classification": TaskType.ZERO_SHOT_CLASSIFICATION } return mapping.get(pipeline_tag) if pipeline_tag else None