Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2
3# Load model directly
4tokenizer = AutoTokenizer.from_pretrained(
5 "Arjun24420/BERT-FakeNews-Classification")
6model = AutoModelForSequenceClassification.from_pretrained(
7 "Arjun24420/BERT-FakeNews-Classification")
8
9# Define class labels mapping
10class_mapping = {
11 0: 'half-true',
12 1: 'mostly-true',
13 2: 'false',
14 3: 'true',
15 4: 'barely-true',
16 5: 'pants-fire'
17}
18
19
20def predict(text):
21 # Tokenize the input text and move tensors to the GPU if available
22 inputs = tokenizer(text, padding=True, truncation=True,
23 max_length=512, return_tensors="pt")
24
25 # Get model output (logits)
26 outputs = model(**inputs)
27
28 probs = outputs.logits.softmax(1)
29 # Get the probabilities for each class
30 class_probabilities = {class_mapping[i]: probs[0, i].item()
31 for i in range(probs.shape[1])}
32
33 return class_probabilities
34
35