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

# 1. Compatibility setup for Hugging Face Spaces
warnings.filterwarnings("ignore")
transformers.PretrainedConfig.is_decoder = False
transformers.PretrainedConfig.output_attentions = False
transformers.PretrainedConfig.output_hidden_states = False

# 2. Path to the model
# When uploaded to HF, the 'model' folder should be in the root
CHECKPOINT_PATH = "model"

print(f"Loading model from: {CHECKPOINT_PATH}...")
# Load the model once at startup
model = ATEPC.AspectExtractor(checkpoint=CHECKPOINT_PATH)

def predict_absa(text):
    if not text.strip():
        return "Please enter some text to analyze."
    
    # Run prediction
    result = model.predict(text, print_result=False)
    
    if not result['aspect']:
        return "No aspects found in the input text."
    
    # Format results for display
    output = []
    for aspect, sentiment in zip(result['aspect'], result['sentiment']):
        output.append({
            "Aspect": aspect,
            "Sentiment": sentiment
        })
    
    return output

# 3. Create Gradio Interface
demo = gr.Interface(
    fn=predict_absa,
    inputs=gr.Textbox(
        lines=3, 
        placeholder="Enter a sentence here (e.g., 'The coffee was great but the price was too high.')",
        label="Input Text"
    ),
    outputs=gr.JSON(label="ABSA Results"),
    title="DeBERTa-v3 Aspect Based Sentiment Analysis",
    description="This demo uses a fine-tuned DeBERTa-v3 model to extract aspects and classify their sentiment polarities.",
    examples=[
        ["The food was delicious but the service was extremely slow."],
        ["The battery life of this laptop is amazing, though the screen is a bit dim."],
        ["I love the interface, but the mobile app crashes frequently."]
    ],
    cache_examples=False
)

if __name__ == "__main__":
    demo.launch()