thavva_rag_ai / howtobuildRAG_model.py

Raghavendra Thavva

First RAG model version

2b1caad about 2 years ago

3.33 kB

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)

	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)