import gradio as gr
import torch
from transformers import pipeline
import os
from huggingface_hub import CommitScheduler
from pathlib import Path
import uuid
import json
import logging

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("app.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("darija-masked-lm")

key = os.environ["HF_KEY"]

submit_file = Path("user_submit/") / f"data_{uuid.uuid4()}.json"
feedback_file = submit_file
submit_file.parent.mkdir(exist_ok=True, parents=True)

scheduler = CommitScheduler(
    repo_id="atlasia/atlaset_inference_ds",
    repo_type="dataset",
    folder_path=submit_file.parent,
    path_in_repo="masked_lm",
    every=5,
    token=key
)

def save_feedback(input, output):
    with scheduler.lock:
        try:
            with feedback_file.open("a") as f:
                f.write(json.dumps({"input": input, "output": output}))
                f.write("\n")
        except Exception as e:
            logger.error(f"Error saving feedback: {str(e)}")

def load_model():
    print("[INFO] Loading model...")
    pipe = pipeline(
        task="feature-extraction",
        model="aitouiazzaneali49/result_model",
        token=key,
        device=-1,
        dtype=torch.float32
    )
    print("[INFO] Model loaded!")
    return pipe

print("[INFO] load model ...")
pipe = load_model()
print("[INFO] model loaded")

def predict(text1, text2):
    emb1 = pipe(text1)[0][0]
    emb2 = pipe(text2)[0][0]
    t1 = torch.tensor(emb1)
    t2 = torch.tensor(emb2)
    similarity = torch.nn.functional.cosine_similarity(t1, t2, dim=0).item()
    result = {"Similarite": round(similarity, 4)}
    save_feedback(f"{text1} | {text2}", result)
    return result

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_text1 = gr.Textbox(label="Phrase 1")
            input_text2 = gr.Textbox(label="Phrase 2")
            submit_btn = gr.Button("Comparer", variant="primary")
            output_labels = gr.Label(label="Résultat")
            submit_btn.click(predict, inputs=[input_text1, input_text2], outputs=output_labels)

demo.queue()
demo.launch()