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_length=50, num_return_sequences=5)
    queries = [llm_tokenizer.decode(seq, 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}"

iface = gr.Interface(
    fn=inpars_v2,
    inputs=gr.Textbox(lines=5, label="Input Document"),
    outputs=gr.Textbox(label="Result"),
    title="InPars-v2 Demo",
    description="Generate queries and train a retriever using LLMs and rerankers."
)

iface.launch()