from datasets import load_dataset
from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer
import faiss
from transformers import DPRContextEncoder, DPRContextEncoderTokenizer
from transformers import BartForConditionalGeneration, BartTokenizer
from transformers import RagTokenForGeneration, RagTokenizer

# Load a Wikipedia dataset as an example
dataset = load_dataset("wikipedia", "20220301.en", split="train")



# Load pre-trained DPR models for retrieval
dpr_question_model = DPRQuestionEncoder.from_pretrained('facebook/dpr-question_encoder-single-nq-base')
dpr_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained('facebook/dpr-question_encoder-single-nq-base')

# Example query embedding
query = "Who won the Nobel Prize in Physics in 2020?"
inputs = dpr_tokenizer(query, return_tensors="pt")
query_embedding = dpr_question_model(**inputs).pooler_output





# Load DPR context encoder
dpr_context_model = DPRContextEncoder.from_pretrained('facebook/dpr-ctx_encoder-single-nq-base')
dpr_context_tokenizer = DPRContextEncoderTokenizer.from_pretrained('facebook/dpr-ctx_encoder-single-nq-base')

# Embed the documents and add them to a FAISS index
document_embeddings = []
faiss_index = faiss.IndexFlatL2(768)  # Assuming 768 is the dimensionality of DPR embeddings

for doc in dataset:
    inputs = dpr_context_tokenizer(doc['text'], return_tensors="pt", truncation=True, max_length=512)
    doc_embedding = dpr_context_model(**inputs).pooler_output.detach().numpy()
    document_embeddings.append(doc_embedding)
    faiss_index.add(doc_embedding)

# Save the FAISS index for later use
faiss.write_index(faiss_index, "document_index.faiss")



# Perform retrieval: Find top-k most similar documents for the query
k = 5  # Number of documents to retrieve
distances, indices = faiss_index.search(query_embedding.detach().numpy(), k)

# Fetch the corresponding documents
retrieved_docs = [dataset[i] for i in indices[0]]
for doc in retrieved_docs:
    print(doc['text'])




# Load a pre-trained BART model
generator_model = BartForConditionalGeneration.from_pretrained('facebook/bart-large')
generator_tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')

# Concatenate the query with the retrieved documents
context = " ".join([doc['text'] for doc in retrieved_docs])
input_text = query + " " + context

# Tokenize the input for BART
inputs = generator_tokenizer(input_text, return_tensors="pt", truncation=True, max_length=1024)

# Generate the response
outputs = generator_model.generate(**inputs, max_length=150, num_beams=5, early_stopping=True)
response = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)

print("Generated Response:", response)





# Load a pre-trained RAG-Token model
rag_model = RagTokenForGeneration.from_pretrained('facebook/rag-token-base')
rag_tokenizer = RagTokenizer.from_pretrained('facebook/rag-token-base')

# Tokenize input
inputs = rag_tokenizer(query, return_tensors="pt", return_token_type_ids=False)
docs_dict = rag_model.retriever(inputs.input_ids, inputs.attention_mask, return_tensors="pt")

# Generate response
generated_ids = rag_model.generate(input_ids=inputs["input_ids"], context_input_ids=docs_dict["context_input_ids"])
generated_text = rag_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

print(generated_text)