entailment.py

from transformers import pipeline

def analyze_entailment(original_sentence, paraphrased_sentences, threshold):
    # Load the entailment model once
    entailment_pipe = pipeline("text-classification", model="ynie/roberta-large-snli_mnli_fever_anli_R1_R2_R3-nli")

    all_sentences = {}
    selected_sentences = {}
    discarded_sentences = {}

    # Prepare input for entailment checks
    inputs = [f"{original_sentence} [SEP] {paraphrase}" for paraphrase in paraphrased_sentences]
    
    # Perform entailment checks for all paraphrased sentences in one go
    entailment_results = entailment_pipe(inputs, return_all_scores=True)

    # Iterate over results
    for paraphrased_sentence, results in zip(paraphrased_sentences, entailment_results):
        # Extract the entailment score for each paraphrased sentence
        entailment_score = next((result['score'] for result in results if result['label'] == 'entailment'), 0)

        all_sentences[paraphrased_sentence] = entailment_score
        
        # Store sentences based on the threshold
        if entailment_score >= threshold:
            selected_sentences[paraphrased_sentence] = entailment_score
        else:
            discarded_sentences[paraphrased_sentence] = entailment_score

    return all_sentences, selected_sentences, discarded_sentences

# Example usage
# print(analyze_entailment("I love you", ["I adore you", "I hate you"], 0.7))