import os
import transformers
from pyabsa import AspectTermExtraction as ATEPC
import warnings

def main():
    # 1. Compatibility setup
    warnings.filterwarnings("ignore")
    transformers.PretrainedConfig.is_decoder = False
    transformers.PretrainedConfig.output_attentions = False
    transformers.PretrainedConfig.output_hidden_states = False

    # 2. Point to your BEST trained model
    CHECKPOINT_PATH = "model"
    
    if not os.path.exists(CHECKPOINT_PATH):
        print(f"Error: Checkpoint not found at {CHECKPOINT_PATH}")
        return

    print(f"Loading model from: {CHECKPOINT_PATH}...")
    
    # 3. Load the model
    # Note: load_model handles the setup automatically
    model = ATEPC.AspectExtractor(checkpoint=CHECKPOINT_PATH)

    print("\n" + "="*50)
    print("READY TO USE ATEPC MODEL")
    print("="*50)
    print("Type a sentence to extract aspects and sentiments.")
    print("Type 'exit' or 'quit' to stop.\n")

    while True:
        text = input("Input text: ").strip()
        
        if text.lower() in ['exit', 'quit', '']:
            break
            
        # Run prediction
        result = model.predict(text, print_result=False)
        
        print(f"Results for: \"{text}\"")
        if not result['aspect']:
            print("  - No aspects found.")
        else:
            for aspect, sentiment in zip(result['aspect'], result['sentiment']):
                print(f"  -> [{aspect:^12}] | Sentiment: {sentiment}")
        print("-" * 30)

if __name__ == "__main__":
    main()