Load a model and store in registry.
(model_name: str)
| 58 | |
| 59 | |
| 60 | def load_model(model_name: str): |
| 61 | """Load a model and store in registry.""" |
| 62 | if model_name in MODEL_REGISTRY: |
| 63 | return MODEL_REGISTRY[model_name] |
| 64 | |
| 65 | if model_name not in MODEL_CONFIGS: |
| 66 | available = list(MODEL_CONFIGS.keys()) |
| 67 | raise ValueError(f"Unknown model '{model_name}'. Available: {available}") |
| 68 | |
| 69 | from funasr import AutoModel |
| 70 | |
| 71 | cfg = MODEL_CONFIGS[model_name].copy() |
| 72 | cfg["device"] = DEVICE |
| 73 | cfg["disable_update"] = True |
| 74 | |
| 75 | logger.info(f"Loading model '{model_name}' on {DEVICE}...") |
| 76 | t0 = time.time() |
| 77 | model = AutoModel(**cfg) |
| 78 | elapsed = time.time() - t0 |
| 79 | logger.info(f"Model '{model_name}' loaded in {elapsed:.1f}s") |
| 80 | |
| 81 | MODEL_REGISTRY[model_name] = model |
| 82 | return model |
| 83 | |
| 84 | |
| 85 | def clean_text(text: str) -> str: |
no test coverage detected
searching dependent graphs…