import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForMaskedLM

# Завантаження додатково навченої моделі
model_path = 'POMAHSLS/BERT_Finetunned_LawTextes'  # Переконайтеся, що шлях до моделі правильний
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForMaskedLM.from_pretrained(model_path)
    
# Ініціалізація пайплайну для завдання fill-mask
fill_mask_pipeline = pipeline("fill-mask", model=model, tokenizer=tokenizer)

# Функція для передбачень
def predict_masked_sentence_with_replacement(sentence):
    """
    Передбачає токени для маскованого слова та замінює [MASK] на найкращий токен.
    Аргументи:
        sentence (str): Текст із токеном [MASK]
    Повертає:
        str: Повне речення із заміною та список передбачень з імовірностями
    """
    if "[MASK]" not in sentence:
        return "Будь ласка, додайте [MASK] у речення для передбачення."
    
    # Отримуємо передбачення
    predictions = fill_mask_pipeline(sentence)
    top_prediction = predictions[0]  # Беремо найкращий результат
    
    # Замінюємо [MASK] на найкращий токен
    predicted_sentence = sentence.replace("[MASK]", top_prediction["token_str"])
    
    # Формуємо рядок для виводу
    prediction_list = "\n".join([
        f"{pred['token_str']} (score: {pred['score']:.4f})"
        for pred in predictions
    ])
    
    # Повертаємо підсумкове речення та список передбачень
    return f"Передбачене речення:\n{predicted_sentence}\n\nМожливі слова:\n{prediction_list}"

# Створення інтерфейсу Gradio
interface = gr.Interface(
    fn=predict_masked_sentence_with_replacement,  # Функція обробки
    inputs=gr.Textbox(
        lines=3, 
        placeholder="Введіть речення з [MASK], наприклад:\nThe [MASK] is responsible for interpreting the Constitution."
    ),
    outputs="text",
    title="BERT_LAW",
    description="Введіть текст з токеном [MASK], і модель замінить його на кращий варіант.",
    examples=[
        ["The [MASK] is responsible for interpreting the Constitution."],
        ["A valid contract requires an offer, [MASK], and consideration."],
        ["Copyright infringement occurs when [MASK] works are copied without permission."]
    ]
)

# Запуск застосунку
interface.launch()