import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
from sentence_transformers import SentenceTransformer
import numpy as np

# Load models
llm = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
llm_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
reranker = AutoModelForSequenceClassification.from_pretrained("cross-encoder/ms-marco-MiniLM-L-6-v2")
reranker_tokenizer = AutoTokenizer.from_pretrained("cross-encoder/ms-marco-MiniLM-L-6-v2")
retriever = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

def generate_query(document):
    prompt = f"Generate a relevant search query for the following document:\n\n{document}\n\nQuery:"
    input_ids = llm_tokenizer.encode(prompt, return_tensors="pt")
    output = llm.generate(
        input_ids, 
        max_new_tokens=30,
        num_return_sequences=5,
        num_beams=5,
        no_repeat_ngram_size=2,
        early_stopping=True
    )
    queries = [llm_tokenizer.decode(seq[input_ids.shape[1]:], skip_special_tokens=True) for seq in output]
    return queries

def rerank_pairs(queries, document):
    pairs = [[query, document] for query in queries]
    inputs = reranker_tokenizer(pairs, padding=True, truncation=True, return_tensors="pt")
    scores = reranker(**inputs).logits.squeeze(-1)
    best_query = queries[torch.argmax(scores)]
    return best_query

def train_retriever(query_doc_pairs):
    # This is a placeholder for the actual training process
    queries, docs = zip(*query_doc_pairs)
    query_embeddings = retriever.encode(queries)
    doc_embeddings = retriever.encode(docs)
    similarity = np.dot(query_embeddings, doc_embeddings.T)
    return f"Retriever trained on {len(query_doc_pairs)} pairs. Average similarity: {similarity.mean():.4f}"

def inpars_v2(document):
    queries = generate_query(document)
    best_query = rerank_pairs(queries, document)
    result = train_retriever([(best_query, document)])
    return f"Generated query: {best_query}\n\n{result}"

# Markdown description of the InPars-v2 paper
paper_description = """
# InPars-v2: Large Language Models as Efficient Dataset Generators for Information Retrieval

**Abstract Link:** [https://arxiv.org/abs/2301.01820](https://arxiv.org/abs/2301.01820)
**PDF Link:** [https://arxiv.org/pdf/2301.01820](https://arxiv.org/pdf/2301.01820)

**Authors:** Vitor Jeronymo, Luiz Bonifacio, Hugo Abonizio, Marzieh Fadaee, Roberto Lotufo, Jakub Zavrel, Rodrigo Nogueira

**Publication Date:** 26 May 2023

## Abstract

Recently, InPars introduced a method to efficiently use large language models (LLMs) in information retrieval tasks: via few-shot examples, an LLM is induced to generate relevant queries for documents. These synthetic query-document pairs can then be used to train a retriever. However, InPars and, more recently, Promptagator, rely on proprietary LLMs such as GPT-3 and FLAN to generate such datasets. In this work we introduce InPars-v2, a dataset generator that uses open-source LLMs and existing powerful rerankers to select synthetic query-document pairs for training. A simple BM25 retrieval pipeline followed by a monoT5 reranker finetuned on InPars-v2 data achieves new state-of-the-art results on the BEIR benchmark. To allow researchers to further improve our method, we open source the code, synthetic data, and finetuned models: [https://github.com/zetaalphavector/inPars/tree/master/tpu](https://github.com/zetaalphavector/inPars/tree/master/tpu)

## Key Features of InPars-v2

1. Uses open-source LLMs for query generation
2. Employs powerful rerankers to select high-quality synthetic query-document pairs
3. Achieves state-of-the-art results on the BEIR benchmark
4. Provides open-source code, synthetic data, and finetuned models

This demo provides a simplified implementation of the InPars-v2 concept, showcasing query generation, reranking, and retriever training.
"""

iface = gr.Interface(
    fn=inpars_v2,
    inputs=gr.Textbox(lines=5, label="Input Document"),
    outputs=gr.Textbox(label="Result"),
    title="InPars-v2 Demo",
    description=paper_description,
    article="This is a minimal implementation of the InPars-v2 concept. For the full implementation and more details, please refer to the original paper and GitHub repository."
)

iface.launch()