from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset
import torch

# Load dataset
dataset = load_dataset('glue', 'mrpc')

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)

# Tokenize dataset
def tokenize_function(examples):
    return tokenizer(examples['sentence1'], examples['sentence2'], truncation=True, padding='max_length', max_length=128)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

# Set up Trainer
training_args = TrainingArguments(
    output_dir='./results',
    evaluation_strategy='epoch',
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
    eval_dataset=tokenized_datasets['validation'],
)

# Train the model
trainer.train()

# Save the model
model.save_pretrained('./my_model')
tokenizer.save_pretrained('./my_model')

# Load the trained model
model = AutoModelForSequenceClassification.from_pretrained('./my_model')
tokenizer = AutoTokenizer.from_pretrained('./my_model')

# Function to predict sentiment
def predict(text1, text2):
    inputs = tokenizer(text1, text2, return_tensors='pt', truncation=True, padding='max_length', max_length=128)
    outputs = model(**inputs)
    logits = outputs.logits
    predictions = torch.argmax(logits, dim=-1).item()
    return "Similar" if predictions == 0 else "Not Similar"

# Web interface
if __name__ == "__main__":
    import gradio as gr

    iface = gr.Interface(
        fn=predict,
        inputs=["text", "text"],
        outputs="text",
        title="MRPC Sentiment Analysis",
        description="Determine if two sentences are semantically equivalent."
    )

    iface.launch()