import gradio as gr
import torch
from model import SimpleMultilingualClassifier  # Import your model

# --- Configuration ---
embedding_files = {
    'en': 'fasttext_embeddings/cc.en.100.bin',
    'fr': 'fasttext_embeddings/cc.fr.100.bin'
    # Add more languages as needed
}
num_classes = 3  # Replace with the actual number of classes
class_labels = ["positive", "negative", "neutral"] # Replace with your actual class labels

# Load the model
try:
    model = SimpleMultilingualClassifier(embedding_files, num_classes)
    # In a real scenario, you would load trained weights here:
    # model.load_state_dict(torch.load('path/to/your/trained_weights.pth'))
    model.eval()
except Exception as e:
    print(f"Error loading model: {e}")
    model = None

def classify_text(text, language):
    if model:
        try:
            prediction = model.predict(text, language, class_labels)
            return prediction
        except ValueError as e:
            return str(e)
    else:
        return "Model not loaded."

iface = gr.Interface(
    fn=classify_text,
    inputs=[
        gr.Textbox(label="Enter text"),
        gr.Dropdown(choices=list(embedding_files.keys()), label="Language")
    ],
    outputs=gr.Textbox(label="Prediction"),
    title="Simple Multilingual Text Classifier",
    description="A basic multilingual text classifier using FastText embeddings.",
)

iface.launch()