get_model_and_inputs() — pytorch Function Reference

Architecture documentation for the get_model_and_inputs() function in huggingface_llm_models.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration calls 1

Entity Profile

DynamoBenchmarks→ Configuration→ get_model_and_inputs() — pytorch Function Reference

Dependency Diagram

graph TD
  a6b96554_3d44_2142_1923_87a603cee1cf["get_model_and_inputs()"]
  6c4a7daf_e704_6d7a_bcaa_98900e3a377b["get_model_and_inputs()"]
  a6b96554_3d44_2142_1923_87a603cee1cf -->|calls| 6c4a7daf_e704_6d7a_bcaa_98900e3a377b
  style a6b96554_3d44_2142_1923_87a603cee1cf fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/huggingface_llm_models.py lines 71–93

    def get_model_and_inputs(model_name, device):
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device)
        model.eval()

        model.generation_config.do_sample = False
        model.generation_config.use_cache = True
        model.generation_config.cache_implementation = "static"
        model.generation_config.max_new_tokens = TextGenerationBenchmark.OUTPUT_LENGTH
        model.generation_config.pad_token_id = tokenizer.eos_token_id
        model.generation_config.temperature = 0.0

        vocab_size = tokenizer.vocab_size
        input_ids = torch.randint(
            low=0,
            high=vocab_size,
            size=(1, TextGenerationBenchmark.INPUT_LENGTH),
            device=device,
            dtype=torch.long,
        )
        example_inputs = {"input_ids": input_ids}

        return model, example_inputs