import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel, PeftConfig

# Load PEFT adapter config
adapter_name = "shukdevdatta123/twitter-distilbert-base-uncased-sentiment-analysis-lora-text-classification"
config = PeftConfig.from_pretrained(adapter_name)

# Load base model and tokenizer
base_model = AutoModelForSequenceClassification.from_pretrained(config.base_model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

# Load the LoRA adapter into the base model
model = PeftModel.from_pretrained(base_model, adapter_name)

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# Label mapping
id2label = {
    0: "Negative",
    1: "Positive"
}

# Prediction function
def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
    
    with torch.no_grad():
        logits = model(**inputs).logits

    predicted_class = torch.argmax(logits, dim=1).item()
    label = id2label[predicted_class]

    probs = torch.nn.functional.softmax(logits, dim=1)
    confidence = probs[0][predicted_class].item()

    return f"{label} (Confidence: {confidence:.2f})"

# Gradio UI
interface = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence to analyze sentiment..."),
    outputs="text",
    title="Twitter Sentiment Classifier (LoRA + DistilBERT)",
    description="This app uses a DistilBERT model with LoRA adapters to classify tweet sentiment as Positive or Negative."
)

interface.launch()