Function run_predict

codegeex/mindspore/generation.py:195–246

run predict

(model_predict, config, args_opt, rank)

def run_predict(model_predict, config, args_opt, rank):
    """run predict"""
    from src.generate import generate, 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",
    ]
    verbose = False
    for i, sample in enumerate(samples):
        for _ in range(1):
            tokenized_token = tokenizer.encode_code(sample)
            input_ids = np.array(tokenized_token).reshape(1, -1)
            # Call inference
            generate_func = generate_increment if config.use_past else generate
            t0 = time.perf_counter()
            output_ids = generate_func(model_predict, input_ids, args_opt, verbose)
            # Decode output ids to sentence
            t1 = time.perf_counter()
            output_samples = tokenizer.decode_code(output_ids.tolist())
            output_samples_str = "".join(output_samples)
            if rank % 8 == 0:
                print(f"=================== prompt {i} ====================")
                print(sample, flush=True)
                print(f"=================== generation {i} ====================")
                print(output_samples_str, flush=True)
                print(
                    f"=== Total time (s): {t1 - t0}, {output_ids.shape[-1] - input_ids.shape[-1]} tokens, {(output_ids.shape[-1] - input_ids.shape[-1]) / (t1 - t0)} token/s")


def main():