import json
import os


def yield_keys_prompts(jsonl_file):
    with open(jsonl_file, 'r') as f:
        for line in f:
            data = json.loads(line)
            #print(data.keys())
            entry = data['entry']
            #print(entry.keys())
            key = entry['index']
            if "result" in data:
                prompts = [data['result']]
                #print(prompts)
            else:
                results = [data['result1'], data['result2']] if 'result2' in entry else [data['result1']]
                prompts = []
                for result in results:
                    prompt = result['extended'] if result.get('extended', None) else result['generated']
                    prompts.append(prompt)
            print(f"Key: {key}, Prompt: {prompts}")
            yield key, prompts
