inference with a model trained on query well-formedness

query_wellformedness_score.py

"""
inference with a model trained on query well-formedness

https://huggingface.co/Ashishkr/query_wellformedness_score

pip transformers install accelerate optimum -q
"""
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Step 1: Initialize tokenizer and model globally
tokenizer = AutoTokenizer.from_pretrained("Ashishkr/query_wellformedness_score")
model = AutoModelForSequenceClassification.from_pretrained(
    "Ashishkr/query_wellformedness_score"
)
model = model.to_bettertransformer()
model.eval()


def get_best_sentence(sentences):
    """
    Takes a list of sentences and returns the sentence with the highest score.

    Args:
    sentences (list): A list of sentences to evaluate.

    Returns:
    str: The sentence with the highest score.
    """

    if isinstance(sentences, str):
        sentences = [sentences]

    # Return the sentence directly if there's only one
    if len(sentences) == 1:
        return sentences[0]

    features = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")

    with torch.no_grad():
        scores = model(**features).logits

        # Check if the output logits have only one dimension
        if scores.shape[1] == 1:
            # If there's only one dimension, use it directly
            well_formed_scores = scores.squeeze()
        else:
            # If there are two dimensions, use softmax to get probabilities
            probabilities = torch.nn.functional.softmax(scores, dim=1)
            well_formed_scores = probabilities[:, 1]

    max_score_idx = torch.argmax(well_formed_scores).item()
    return sentences[max_score_idx]


# Example usage
sentences = [
    "The quarterly financial report are showing an increase.",
    "Him has completed the audit for last fiscal year.",
    "Please to inform the board about the recent developments.",
    "The team successfully achieved all its targets for the last quarter.",
    "Our company is exploring new ventures in the European market.",
]

best_sentence = get_best_sentence(sentences)
print(f"The best sentence is: '{best_sentence}'")