import gradio as gr
import torch
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification, AutoConfig
import os

model_name = "distilbert-base-uncased"
model_path = "cefr_model.pth"
num_labels = 6
level_mapping = {0: 'A1', 1: 'A2', 2: 'B1', 3: 'B2', 4: 'C1', 5: 'C2'}

tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
config = AutoConfig.from_pretrained(model_name, num_labels=num_labels)
model = DistilBertForSequenceClassification.from_pretrained(model_name, config=config)
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval()

def predict_cefr(text):
    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
        predicted_class = torch.argmax(probabilities).item()
    return level_mapping[predicted_class]

iface = gr.Interface(
    fn=predict_cefr,
    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
    outputs="text",
    title="CEFR Level Prediction",
    description="Enter some English text and I'll predict its CEFR level!"
)

iface.launch()