Create app.py

app.py

@@ -0,0 +1,50 @@
+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()