classify.py

import src.preparemodel
import torch

class RumourDetectClass:
    def __init__(self):
        self.model, self.tokenizer = src.preparemodel.getmodel(True)
        self.max_length = 128
    def classify(self, text: str) -> int:
        self.model.to('cpu')
        self.model.eval()
        encoding = self.tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt"
        )
        input_ids = encoding["input_ids"].flatten().unsqueeze(0)
        attention_mask = encoding["attention_mask"].flatten().unsqueeze(0)
        outputs = self.model(input_ids, attention_mask=attention_mask)
        pred_class = torch.argmax(outputs, dim=-1).item()
        return pred_class
    def __call__(self, text: str) -> int:
        return self.classify(text=text)

if __name__ == "__main__":
    import pandas as pd
    csv_file = 'data/val.csv'
    data = pd.read_csv(csv_file)
    texts = data['text'].tolist()
    labels = data['label'].tolist()

    test_idx = 0

    classifier = RumourDetectClass()

    pred = classifier(texts[test_idx])
    print(f"predict\t{pred}\nlabel\t{labels[test_idx]}")