# Author: Siwen Yu (yusiwen@gmail.com)
# License: MIT

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel, PeftConfig

# 1. Define paths
base_model_path = "bert-base-uncased"
lora_weights_path = "./my_saved_lora"

# 2. Load tokenizer and the clean, base model
tokenizer = AutoTokenizer.from_pretrained(lora_weights_path)
base_model = AutoModelForSequenceClassification.from_pretrained(base_model_path, num_labels=2)

# 3. Mathematically merge the LoRA weights on top of the base model
model = PeftModel.from_pretrained(base_model, lora_weights_path)
model = model.to("cuda")
model.eval() # Set to evaluation mode

# 4. Run custom evaluation samples
test_sentences = [
    "This movie was an absolute masterpiece with breathtaking visuals!", # Expect Positive (1)
    "A total waste of time. The acting was horrible and the plot made no sense." # Expect Negative (0)
]

print("\n--- Running Validation Inference ---")
for text in test_sentences:
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128).to("cuda")
    
    with torch.no_grad():
        outputs = model(**inputs)
        prediction = torch.argmax(outputs.logits, dim=-1).item()
        
    sentiment = "Positive" if prediction == 1 else "Negative"
    print(f"Review: '{text}' -> Predicted Sentiment: {sentiment}")