Function run_predict

codegeex/mindspore/generation_1p.py:201–239

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 = [
        "# 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: 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",
        "bool prime(int n) {\n    // Find whether n is a prime number\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")
        break


def main():