def main()

in ultravox/tools/push_to_hub.py [0:0]


def main(args: UploadToHubArgs):
    # Load the model and tokenizer, then merge LoRA weights if they exist
    inference = ultravox_infer.UltravoxInference(
        args.model,
        tokenizer_id=args.tokenizer,
        device=args.device,
        data_type=args.data_type,
    )
    pipe = ultravox_pipeline.UltravoxPipeline(
        model=inference.model,
        tokenizer=inference.tokenizer,
        audio_processor=inference.processor.audio_processor,
        device=args.device,
    )
    print("Uploading model to HuggingFace Hub...")
    pipe.push_to_hub(args.hf_upload_model, private=args.private)

    if args.verify:
        print("Model uploaded. Testing model...")
        loaded_pipe = transformers.pipeline(
            model=args.hf_upload_model, trust_remote_code=True
        )
        ds = datasets.create_dataset("boolq", datasets.VoiceDatasetArgs())
        sample = next(iter(ds))
        generated = loaded_pipe(
            {"audio": sample.audio, "turns": sample.messages[:-1]}, max_new_tokens=10
        )
        print(f"Generated (max 10 tokens): {generated}")