Spaces:

KeisukeMiyamoto
/

lambda-audio-v6

Sleeping

lambda-audio-v6 / src /modeling /checkpoint.py

Add simple LambdaAudio v6 ZeroGPU transcription demo

82e11f0 verified 22 days ago

577 Bytes

from pathlib import Path
from typing import Any, cast

import torch
from torch import nn


def load_checkpoint_state(path: Path) -> dict[str, torch.Tensor]:
    checkpoint = cast(
        dict[str, Any],
        torch.load(path, map_location="cpu", weights_only=True),
    )
    return cast(dict[str, torch.Tensor], checkpoint["state_dict"])


def load_model_checkpoint(model: nn.Module, path: Path) -> None:
    state_dict = {name.removeprefix("model."): parameter for name, parameter in load_checkpoint_state(path).items()}
    model.load_state_dict(state_dict, strict=True)

	from pathlib import Path
	from typing import Any, cast

	import torch
	from torch import nn


	def load_checkpoint_state(path: Path) -> dict[str, torch.Tensor]:
	checkpoint = cast(
	dict[str, Any],
	torch.load(path, map_location="cpu", weights_only=True),
	)
	return cast(dict[str, torch.Tensor], checkpoint["state_dict"])


	def load_model_checkpoint(model: nn.Module, path: Path) -> None:
	state_dict = {name.removeprefix("model."): parameter for name, parameter in load_checkpoint_state(path).items()}
	model.load_state_dict(state_dict, strict=True)