import json
import numpy as np

with open("./retrieval/data/inauto_examples.json", "r") as file:  
    prob_message = json.load(file)
prob_message = prob_message['messages'] 

def prob_sample(input, k=8):
    indices = np.arange(int((len(prob_message)-1)/2))
    np.random.shuffle(indices)
    prob_example = [prob_message[0]]
    for (n,ind) in enumerate(indices):
        if n > k:
            break
        else:
            prob_example += prob_message[2*ind+1:2*ind+3]
    return prob_example
