Method run_test

scripts/eval_needle.py:209–292

(self)

        return int(context_length)

    def run_test(self):
        fs = gcsfs.GCSFileSystem()
        contexts = []
        template = self.OURS_TEMPLATE

        def _key_from_result(result):
            return (result['context_length'], result['depth_percent'], result['seed'])

        results = []
        completed = set()
        def exists(fname):
            if fname.startswith('gs://'):
                return fs.exists(fname)
            else:
                return os.path.exists(fname)
        if exists(FLAGS.output_file):
            with open_file(FLAGS.output_file, 'r') as f:
                results = json.load(f)
                completed = set([_key_from_result(result) for result in results])
        print('completed', len(completed))

        full_contexts = self.read_context_files(FLAGS.n_rounds)
        full_tokens = [self.enc.encode(full_context) for full_context in tqdm(full_contexts)]

        start = time.time()
        for context_length in self.context_lengths:
            trim_contexts = [self.enc.decode(full_token[:context_length]) for full_token in tqdm(full_tokens)]
            max_input_length = self.compute_max_input_length(context_length)
            contexts = []
            for depth_percent in self.document_depth_percents:
                for i in range(FLAGS.n_rounds):
                    if (int(context_length), float(depth_percent), i) in completed:
                        continue
                    random_city = random.choice(LLMNeedleHaystackTester.RANDOM_NEEDLE_CITIES)
                    insert_needle = True
                    needle_rnd_number = str(self.generate_random_number(self.rnd_number_digits))
                    print("context length: " + str(context_length))
                    print("depth_percent : " + str(depth_percent))
                    context = self.create_contexts(needle_rnd_number, insert_needle, random_city, trim_contexts[i], context_length, depth_percent, i)
                    contexts.append(context)

            if len(contexts) == 0:
                continue

            B = FLAGS.max_tokens_per_batch / (max_input_length + self.model.block_size)
            B = int(B / self.model.data_dim) * self.model.data_dim
            if B < self.model.data_dim:
                B = self.model.data_dim
            elif B > len(contexts):
                B = int(math.ceil(len(contexts) / self.model.data_dim) * self.model.data_dim)
            if len(contexts) % B == 0:
                n_pad = 0
            else:
                n_pad = B - len(contexts) % B
            for _ in range(n_pad):
                contexts.insert(0, contexts[0])

            pbar = tqdm(total=len(contexts))