from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# Load the model and tokenizer from Hugging Face Hub
model_name = "vinD27/stock_news"  # Replace with your model repo name
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Map label indices to human-readable class names
label_mapping = {
    0: "negative",
    1: "neutral",
    2: "positive"
}

# Input text
input_text = "Wow. The stock is amazing"

# Tokenize and predict
inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True, max_length=128)
outputs = model(**inputs)
predicted_class_idx = torch.argmax(outputs.logits, dim=-1).item()  # Get the predicted label index

# Print the results
print(f"Your input is: '{input_text}'")
print(f"And the prediction is: {label_mapping[predicted_class_idx]} ({predicted_class_idx})")