Method run_test

scripts/eval_needle_multi.py:214–301

(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 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 full_tokens]
            max_input_length = self.compute_max_input_length(context_length)
            contexts = []
            for i in range(FLAGS.n_rounds):
                if (int(context_length), i) in completed:
                    continue
                random_cities = random.sample(LLMNeedleHaystackTester.RANDOM_NEEDLE_CITIES, FLAGS.n_needles_total)
                document_depths = random.sample(self.document_depth_percents, FLAGS.n_needles_total)
                random_cities_retrieve = random.sample(random_cities, FLAGS.n_needles_retrieve)
                needles_info = {}
                for random_city, depth_percent in zip(random_cities, document_depths):
                    needles_info[random_city] = (
                        str(self.generate_random_number(self.rnd_number_digits)),
                        depth_percent
                    )
                context = self.create_contexts(needles_info, random_cities_retrieve, trim_contexts[i], context_length, 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)
            n_pad = B - len(contexts) % B
            for _ in range(n_pad):
                contexts.insert(0, contexts[0])

            pbar = tqdm(total=len(contexts))