| 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) | |