app.py

import py_vncorenlp
from sentence_transformers import CrossEncoder
py_vncorenlp.download_model(save_dir='/absolute/path/to/vncorenlp')
rdrsegmenter = py_vncorenlp.VnCoreNLP(annotators=["wseg"], save_dir='/absolute/path/to/vncorenlp')

def rerank(query,sentences):
    tokenized_query = rdrsegmenter.word_segment(query)
    tokenized_sentences = [rdrsegmenter.word_segment(sent) for sent in sentences]

    tokenized_pairs = [[tokenized_query, sent] for sent in tokenized_sentences]

    MODEL_ID = 'itdainb/PhoRanker'
    MAX_LENGTH = 512 

    model = CrossEncoder(MODEL_ID, max_length=MAX_LENGTH)
    
    # For fp16 usage
    model.model.half()

    scores = model.predict(tokenized_pairs)

    # 0.982, 0.2444, 0.9253
    'print(scores)'
    return scores




# Create Gradio interface
interface = gr.Interface(
    fn=rerank,
    inputs=[
        gr.Textbox(label="Query", placeholder="Enter your query"),
        gr.Textbox(label="Documents (one per line)", lines=5, placeholder="Enter documents to rank"),
    ],
    outputs=gr.Textbox(label="Reranked Documents"),
    title="MonoT5 Reranking",
    description="Provide a query and a list of documents to rerank them using MonoT5."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()