chatrag-deberta / inference_chatrag.py
Pclanglais's picture
Upload inference_chatrag.py
4a1b022 verified
import pandas as pd
import torch
import numpy as np
from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer
model = AutoModelForSequenceClassification.from_pretrained("deberta-classification-chatrag/checkpoint-6342")
tokenizer = AutoTokenizer.from_pretrained("deberta-classification-chatrag/checkpoint-6342")
result = ["Comment puis-je renouveler un passeport ?", "Combien font deux et deux ?", "Écris un début de lettre de recommandation pour la Dinum"]
result = pd.DataFrame(result, columns=['query'])
complete_probabilities = []
for text in result["query"].tolist():
encoding = tokenizer(text, return_tensors="pt")
encoding = {k: v.to(model.device) for k,v in encoding.items()}
outputs = model(**encoding)
logits = outputs.logits
logits.shape
# apply sigmoid + threshold
sigmoid = torch.nn.Sigmoid()
probs = sigmoid(logits.squeeze().cpu())
predictions = np.zeros(probs.shape)
# Extract the float value from the tensor
float_value = probs.item()
complete_probabilities.append(float_value)
result["prob"] = complete_probabilities
print(result)