Function run_predict

codegeex/mindspore/generation_batch.py:193–295

run predict

(model_predict, config, args_opt, rank)

def run_predict(model_predict, config, args_opt, rank):
    """run predict"""
    from src.generate_finetune import generate_increment
    # Define tokenizer
    tokenizer = CodeTokenizer(mode='6b')

    # Tokenize input sentence to ids
    samples = [
        "Hello there!",
        "# language: Python\ndef add(a, b):\n    '''\n    Find the sum of a and b.\n    '''\n",
        "def add(a, b):\n    '''\n    Find the sum of a and b.\n    '''\n",
        "# language: Python\ndef optimization():\n    '''\n    Find the maximum of P=E**2*R/(R + r)**2 if E and r are fixed but R varies. Import sympy. Use sympy. Find where the derivative is equal to zero. Substitute the value of R into P.\n    '''\n",
        "from typing import List\n\n\ndef has_close_elements(numbers: List[float], threshold: float) -> bool:\n    \"\"\" Check if in given list of numbers, are any two numbers closer to each other than\n    given threshold.\n    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)\n    False\n    >>> has_close_elements([1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)\n    True\n    \"\"\"\n",
        "// language: JavaScript\nfunction prime(n) {\n    // Find whether n is a prime number.\n",
        "string morse_encoder(string text) {\n    // Translate text into Morse code\n",
        "def morse_encoder(text):\n    # Translate text into Morse code separated by spaces\n",
        f"% language: MATLAB\nfunction x = solve(A, b)\n    % Solve Ax = b\n",
        f"% language: MATLAB\nfunction [L, U] = lu(A)\n    % Return LU factorization of A\n",
        "def TCPState(state):\n    # given a state in TCP protocol, return a list of next possible states\n",
        "def coordinates(p1, p2, precision=0)\n    # p1 is (x1, y1), p2 is (x2, y2), return the distance between p1 and p2 on a cartesian plane, rounded to precision\n",
        "double travel(double total_time, double run_time, double rest_time, double speed) {\n    // the horse runs for run_time with speed speed and rests for rest_time, return the distance it travels after total_time\n",
        "def travel(total_time, run_time, rest_time, speed):\n    # the horse runs for run_time with speed speed and rests for rest_time, return the distance it travels after total_time\n",
        "// language: C++\nint add(int a, int b) {\n    /* Find the sum of a and b. */\n",
        "int add(int a, int b) {\n    /* Find the sum of a and b. */\n",
        "// language: C++\nvoid sort(int *array, int len) {\n   // Sort the array with length len\n",
        "bool prime(int n) {\n    // Find whether n is a prime number\n",
        "def prime(n):\n    # Find whether n is a prime number\n",
        f"% language: MATLAB\nfunction H = hilbert(n)\n    % Return Hilbert matrix of size n * n\n",
        f"% language: MATLAB\nfunction L = cholesky(A)\n    % Return Cholesky factorization of symmetric positive definete matrix A\n",
        "// language: JavaScript\nfunction add(a, b) {\n    // Find the sum of a and b.\n",
        "# language: R\nadd<-function(a, b) {\n    # Find the sum of a and b.\n",
    ]
    samples = [tokenizer.encode_code(l) for l in samples]
    generations = []
    batch_size = config.batch_size
    verbose = (rank % 8 == 0)
    save_path = f'/home/work/sfs/xx/pangu_alpha_code/generation_batch/{args_opt.temperature}.txt'  # TODO: set as current save path
    save_dir = os.path.split(save_path)[0]
    if rank == 0:
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        if not os.path.exists(save_path):
            f = open(save_path, 'w')
            f.close()
        os.system(f'sudo chmod 777 -R {save_dir}')
    batch = []
    input_length = []
    sample_ids = []
    for i, sample in enumerate(samples):
        tokenized_token = sample
        input_ids = np.array(tokenized_token).reshape(1, -1)
        batch.append(input_ids)
        input_length.append(input_ids.shape[1])
        sample_ids.append(i)
        if (i + 1) % batch_size == 0:
            valid_length = max(input_length)
            for j in range(len(batch)):
                batch[j] = np.pad(batch[j], ((0, 0), (0, valid_length - input_length[j])),