Spaces:
Configuration error
Configuration error
File size: 7,809 Bytes
af2c3f6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """
Main orchestrator for the HuggingFace Model Selector
"""
import sys
import os
from pathlib import Path
# Add the project root to Python path
project_root = str(Path(__file__).parent.parent.absolute())
if project_root not in sys.path:
sys.path.insert(0, project_root)
import asyncio
from typing import Optional
from src.agents.input_agent import InputAgent
from src.agents.research_agent import ResearchAgent
from src.agents.evaluation_agent import EvaluationAgent
from src.agents.benchmarking_agent import BenchmarkingAgent
from src.agents.deployment_agent import DeploymentAgent
from src.models.schemas import DeploymentType, SelectionResult, TaskType
class HuggingFaceModelSelector:
"""
Main orchestrator that coordinates all agents.
This class:
1. Takes a task description
2. Parses requirements
3. Searches for models
4. Evaluates and scores them
5. Benchmarks top models
6. Generates deployment code
"""
def __init__(self):
self.input_agent = InputAgent()
self.research_agent = ResearchAgent()
self.evaluation_agent = EvaluationAgent()
self.benchmarking_agent = BenchmarkingAgent()
self.deployment_agent = DeploymentAgent()
async def select_and_deploy(self,
task_description: str,
deployment_type: DeploymentType = DeploymentType.FASTAPI,
benchmark: bool = True,
top_k: int = 5) -> SelectionResult:
"""
Complete pipeline: select best model and generate deployment code.
Args:
task_description: Natural language task (e.g., "translate english to french")
deployment_type: Type of deployment to generate
benchmark: Whether to run performance benchmarks
top_k: Number of top models to consider
Returns:
SelectionResult with all details
"""
try:
print("\n" + "=" * 60)
print(" HuggingFace Model Selector")
print("=" * 60)
# Step 1: Parse requirements
print("\n Step 1: Analyzing requirements...")
requirements = self.input_agent.parse_requirements(task_description)
print(f" Task: {requirements.task_type.value}")
# Display task-specific requirements
if requirements.task_type == TaskType.TRANSLATION and requirements.translation_reqs:
req = requirements.translation_reqs
print(f" Translation: {req.source_language.value} → {req.target_language.value}")
if req.domain:
print(f" Domain: {req.domain}")
elif requirements.task_type == TaskType.TEXT_TO_SPEECH and requirements.tts_reqs:
req = requirements.tts_reqs
print(f" TTS Language: {req.language.value}")
print(f" Voice: {req.voice_type.value}")
elif requirements.task_type == TaskType.SPEECH_TO_TEXT and requirements.stt_reqs:
req = requirements.stt_reqs
print(f" STT Language: {req.language.value}")
if req.domain:
print(f" Domain: {req.domain}")
elif requirements.task_type in [TaskType.TEXT_GENERATION, TaskType.CHAT,
TaskType.INSTRUCTION_FOLLOWING, TaskType.CODE_GENERATION,
TaskType.QUESTION_ANSWERING, TaskType.SUMMARIZATION] and requirements.llm_reqs:
req = requirements.llm_reqs
print(f" LLM Size: {req.model_size.value}")
print(f" Context Length: {req.context_length}")
elif requirements.task_type in [TaskType.OCR, TaskType.DOCUMENT_UNDERSTANDING] and requirements.ocr_reqs:
req = requirements.ocr_reqs
langs = [lang.value for lang in req.languages]
print(f" OCR Languages: {langs}")
if req.handwritten:
print(f" Handwriting: Yes")
print(f" Hardware: {[c.value for c in requirements.hardware_constraints]}")
# Step 2: Search for models
print("\n Step 2: Searching HuggingFace...")
models = await self.research_agent.search_models(requirements, top_k=top_k*2)
if not models:
return SelectionResult(
status="error",
error="No models found matching your requirements"
)
# Step 3: Score models
print("\n Step 3: Evaluating models...")
scored_models = self.evaluation_agent.score_models(models, requirements)
# Display top models
print("\n Top Models:")
for i, scored in enumerate(scored_models[:5]):
print(f" {i+1}. {scored.model_id} (Score: {scored.total_score:.3f})")
# Show top 3 component scores
top_metrics = sorted(scored.component_scores.items(), key=lambda x: x[1], reverse=True)[:3]
for metric, score in top_metrics:
print(f" - {metric}: {score:.2f}")
# Step 4: Select best model
best_model = scored_models[0]
print(f"\n Step 4: Selected model: {best_model.model_id}")
# Step 5: Benchmark (optional)
benchmark_results = []
if benchmark:
print("\n Step 5: Running benchmarks...")
benchmark_results = await self.benchmarking_agent.benchmark_models(
[m.model_id for m in scored_models[:3]],
requirements.task_type,
requirements
)
if benchmark_results and not benchmark_results[0].error:
print(f"\n Benchmark Results for {best_model.model_id}:")
print(f" Latency: {benchmark_results[0].latency_ms:.2f} ms")
print(f" Memory: {benchmark_results[0].memory_usage_mb:.2f} MB")
if benchmark_results[0].throughput:
print(f" Throughput: {benchmark_results[0].throughput:.2f} samples/sec")
# Step 6: Generate deployment code
print("\n Step 6: Generating deployment code...")
deployment_files = self.deployment_agent.generate_deployment(
model_id=best_model.model_id,
task_type=requirements.task_type,
deployment_type=deployment_type,
benchmark_results=benchmark_results[0] if benchmark_results else None,
requirements=requirements
)
# Save files
output_folder = self.deployment_agent.save_deployment_files(deployment_files)
print(f"\n Deployment files saved to: {output_folder}")
return SelectionResult(
status="success",
selected_model=best_model.model_id,
task_type=requirements.task_type,
requirements=requirements,
all_scores=scored_models,
benchmark_results=benchmark_results,
deployment_files=deployment_files
)
except Exception as e:
print(f"\n Error: {e}")
import traceback
traceback.print_exc()
return SelectionResult(
status="error",
error=str(e)
)
|