INNER CODE UNIT · Python

main

KodCode-AI/kodcode · pipeline/completion_open_model.py:215

def main():
    # Load instructions from the input file
    dataset = load_dataset_from_file(INPUT_FILE_NAME)
    
    # Initialize the engine
    if args.engine == "together":
        print("Start together API engine...")
        llm = None
        params = None
        tokenizer = None
    elif args.engine == "vllm":
        # Set the device
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device
        print("Start Local vllm engine...")
        llm = LLM(model=MODEL_NAME, 
            dtype=args.dtype,
            trust_remote_code=True,
            max_model_len = args.max_model_len, # limited by kv-cache 

View source record →