import torch from transformers import DistilBertTokenizer, DistilBertModel # Load the tokenizer and model tokenizer = DistilBertTokenizer.from_pretrained("tokenizer_config.json") model = DistilBertModel.from_pretrained("pytorch_model.bin") # Define the inference function def predict(text): # Tokenize the input inputs = tokenizer(text, padding="max_length", truncation=True, return_tensors="pt") # Perform the inference with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits # Convert logits to probabilities probabilities = torch.softmax(logits, dim=1).squeeze().tolist() return probabilities # Example usage text = "This is a sample input." probabilities = predict(text) print(probabilities)