File size: 1,166 Bytes
6582310 fd49f40 6582310 fd49f40 6582310 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
# Load the model and tokenizer
model_name = "TerminatorPower/bert-news-classif-turkish"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()
# Load the reverse label mapping
reverse_label_mapping = {
0: "turkiye",
1: "ekonomi",
2: "dunya",
3: "spor",
4: "magazin",
5: "guncel",
6: "genel",
7: "siyaset",
8: "saglik",
9: "kultur-sanat",
10: "teknoloji",
11: "yasam"
}
def predict(text):
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
inputs = {key: value.to("cuda" if torch.cuda.is_available() else "cpu") for key, value in inputs.items()}
model.to(inputs["input_ids"].device)
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.argmax(outputs.logits, dim=1)
predicted_label = reverse_label_mapping[predictions.item()]
return predicted_label
if __name__ == "__main__":
text = input()
print(f"Predicted label: {predict(text)}")
|