import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from trainml import train_and_save_model  # Import the training function
train_and_save_model()

# Load the trained model and tokenizer
model_path = "path/to/save/model"
tokenizer_path = "path/to/save/tokenizer"

model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
model.eval()  # Set model to evaluation mode

def predict_paraphrase(sentence1, sentence2):
    # Tokenize the input sentences
    inputs = tokenizer(sentence1, sentence2, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Get probabilities
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1).tolist()[0]
    
    # Assuming the first class (index 0) is 'not paraphrase' and the second class (index 1) is 'paraphrase'
    return {"Not Paraphrase": probs[0], "Paraphrase": probs[1]}

# Create Gradio interface
iface = gr.Interface(
    fn=predict_paraphrase,
    inputs=[gr.inputs.Textbox(lines=2, placeholder="Enter Sentence 1 Here..."),
            gr.inputs.Textbox(lines=2, placeholder="Enter Sentence 2 Here...")],
    outputs=gr.outputs.Label(num_top_classes=2),
    title="Paraphrase Identification",
    description="This model predicts whether two sentences are paraphrases of each other."
)

iface.launch()