prepare_snli.py

def prepare_snli(sa, sb, lb):
    
    classes = {"entailment", "contradiction"}
    anc_to_pairs = defaultdict(list)
    filtered = {}
    skipped = 0
    anchor_id = 0

    for xa, xb, y in zip(sa, sb, lb):
        anc_to_pairs[xa].append((xb, y))

    for anchor, payload in anc_to_pairs.items(): 
        
        filtered[anchor_id] = defaultdict(list)
        filtered[anchor_id]["anchor"].append(anchor)
        
        labels = set([t[1] for t in payload])
        
        if len(labels&classes) == len(classes):
            for text, label in payload:
                filtered[anchor_id][label].append(text)
            anchor_id += 1
        else:
            skipped += 1
            
    print("Loaded: {} \nSkipped: {}".format(anchor_id, skipped))
        
    return filtered