import torch
import gradio as gr
from transformers import AutoModel

# 1. Инициализация модели (выполняется один раз при запуске)
device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = 'jinaai/jina-reranker-v3'

print(f"Загрузка модели {model_name} на {device}...")
model = AutoModel.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    trust_remote_code=True,
).to(device)
model.eval()

import json

def rerank_documents(query, docs_text):
    if not query.strip() or not docs_text.strip():
        return json.dumps({"error": "Пожалуйста, введите запрос и хотя бы один документ."}, ensure_ascii=False)

    # Разделяем текст на список документов по строкам
    documents = [d.strip() for d in docs_text.split('\n') if d.strip()]
    
    try:
        results = model.rerank(query, documents)
        
        # Создаем словарь: документ -> скор
        score_dict = {res['document']: float(res['relevance_score']) for res in results}
        
        # Формируем список скор в том же порядке, что и документы
        scores = [score_dict.get(doc, 0.0) for doc in documents]
        
        return json.dumps(scores, ensure_ascii=False, indent=2)
    
    except Exception as e:
        return json.dumps({"error": f"Ошибка при обработке: {str(e)}"}, ensure_ascii=False)

# 2. Создание интерфейса Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"# Jina Reranker v3 ({device})")
    gr.Markdown("Введите поисковый запрос и список документов (каждый с новой строки) для оценки релевантности.")
    
    with gr.Row():
        with gr.Column():
            query_input = gr.Textbox(
                label="Запрос (Query)", 
                placeholder="Например: beach running",
                lines=1
            )
            docs_input = gr.Textbox(
                label="Документы (по одному на строку)", 
                placeholder="dress\nrunning shoes\nswimming suit",
                lines=10
            )
            submit_btn = gr.Button("Ранжировать", variant="primary")
        
        with gr.Column():
            output_display = gr.Textbox(
                label="Результаты (Sorted by Score)", 
                lines=15,
                interactive=False
            )

    submit_btn.click(
        fn=rerank_documents, 
        inputs=[query_input, docs_input], 
        outputs=output_display
    )

# 3. Запуск
if __name__ == "__main__":
    demo.launch(share=True)