from llama_index.readers.web import SimpleWebPageReader
 from llama_index.core.node_parser import SentenceSplitter
 
 language = "EN"
 url_doc = "https://eur-lex.europa.eu/legal-content/"+language+"/TXT/HTML/?uri=CELEX:52021PC0206"
 
 documents = SimpleWebPageReader(html_to_text=True).load_data([url_doc])
 
 parser = SentenceSplitter(chunk_size=1000)
 nodes = parser.get_nodes_from_documents(documents, show_progress=True)

prompts={}
 prompts["EN"] = """\
 Context information is below.
 
 ---------------------
 {context_str}
 ---------------------
 
 Given the context information and not prior knowledge, generate only questions based on the below query.
 
 You are a Teacher/ Professor. Your task is to setup {num_questions_per_chunk} questions for an upcoming quiz/examination.
 The questions should be diverse in nature across the document. Restrict the questions to the context information provided."
 """

from llama_index.llms import OpenAI
 from llama_index.legacy.finetuning import generate_qa_embedding_pairs
 
 qa_dataset = generate_qa_embedding_pairs(
    llm=OpenAI(model="gpt-3.5-turbo-0125",additional_kwargs={'seed':42}),
    nodes=nodes,
    qa_generate_prompt_tmpl = prompts[language],
    num_questions_per_chunk=2
 )

def evaluate(dataset, embed_model, insert_batch_size=1000, top_k=5):
    # Get corpus, queries, and relevant documents from the qa_dataset object
    corpus = dataset.corpus
    queries = dataset.queries
    relevant_docs = dataset.relevant_docs
 
    # Create TextNode objects for each document in the corpus and create a VectorStoreIndex to efficiently store and retrieve embeddings
    nodes = [TextNode(id_=id_, text=text) for id_, text in corpus.items()]
    index = VectorStoreIndex(
        nodes, embed_model=embed_model, insert_batch_size=insert_batch_size
    )
    retriever = index.as_retriever(similarity_top_k=top_k)
 
    # Prepare to collect evaluation results
    eval_results = []
 
    # Iterate over each query in the dataset to evaluate retrieval performance
    for query_id, query in tqdm(queries.items()):
        # Retrieve the top_k most similar documents for the current query and extract the IDs of the retrieved documents
        retrieved_nodes = retriever.retrieve(query)
        retrieved_ids = [node.node.node_id for node in retrieved_nodes]
 
        # Check if the expected document was among the retrieved documents
        expected_id = relevant_docs[query_id][0]
        is_hit = expected_id in retrieved_ids # assume 1 relevant doc per query
 
        # Calculate the Mean Reciprocal Rank (MRR) and append to results
        if is_hit:
            rank = retrieved_ids.index(expected_id) + 1
            mrr = 1 / rank
        else:
            mrr = 0
        eval_results.append(mrr)
 
    # Return the average MRR across all queries as the final evaluation metric
    return np.average(eval_results)

from llama_index.embeddings.openai import OpenAIEmbedding
 
 embed_model = OpenAIEmbedding(model=model_spec['model_name'],
                              dimensinotallow=model_spec['dimensions'])

embeddings_model_spec = {
 }
 
 embeddings_model_spec['OAI-Large-256']={'model_name':'text-embedding-3-large','dimensions':256}
 embeddings_model_spec['OAI-Large-3072']={'model_name':'text-embedding-3-large','dimensions':3072}
 embeddings_model_spec['OAI-Small']={'model_name':'text-embedding-3-small','dimensions':1536}
 embeddings_model_spec['OAI-ada-002']={'model_name':'text-embedding-ada-002','dimensions':None}
 
 results = []
 
 languages = ["EN", "FR", "CS", "HU"]
 
 # Loop through all languages
 for language in languages:
 
    # Load dataset
    file_name=language+"_dataset.json"
    qa_dataset = EmbeddingQAFinetuneDataset.from_json(file_name)
 
    # Loop through all models
    for model_name, model_spec in embeddings_model_spec.items():
 
        # Get model
        embed_model = OpenAIEmbedding(model=model_spec['model_name'],
                                      dimensinotallow=model_spec['dimensions'])
 
        # Assess embedding score (in terms of MRR)
        score = evaluate(qa_dataset, embed_model)
 
        results.append([language, model_name, score])
 
 df_results = pd.DataFrame(results, columns = ["Language" ,"Embedding model", "MRR"])

embeddings_model_spec = {
 }
 
 embeddings_model_spec['E5-mistral-7b']={'model_name':'intfloat/e5-mistral-7b-instruct','max_length':32768, 'pooling_type':'last_token', 
                                        'normalize': True, 'batch_size':1, 'kwargs': {'load_in_4bit':True, 'bnb_4bit_compute_dtype':torch.float16}}
 embeddings_model_spec['ML-E5-large']={'model_name':'intfloat/multilingual-e5-large','max_length':512, 'pooling_type':'mean', 
                                      'normalize': True, 'batch_size':1, 'kwargs': {'device_map': 'cuda', 'torch_dtype':torch.float16}}
 embeddings_model_spec['BGE-M3']={'model_name':'BAAI/bge-m3','max_length':8192, 'pooling_type':'cls', 
                                  'normalize': True, 'batch_size':1, 'kwargs': {'device_map': 'cuda', 'torch_dtype':torch.float16}}
 embeddings_model_spec['Nomic-Embed']={'model_name':'nomic-ai/nomic-embed-text-v1','max_length':8192, 'pooling_type':'mean', 
                                      'normalize': True, 'batch_size':1, 'kwargs': {'device_map': 'cuda', 'trust_remote_code' : True}}
 
 results = []
 
 languages = ["EN", "FR", "CS", "HU"]
 
 # Loop through all models
 for model_name, model_spec in embeddings_model_spec.items():
 
    print("Processing model : "+str(model_spec))
 
    # Get model
    tokenizer = AutoTokenizer.from_pretrained(model_spec['model_name'])
    embed_model = AutoModel.from_pretrained(model_spec['model_name'], **model_spec['kwargs'])
         
    if model_name=="Nomic-Embed":
        embed_model.to('cuda')
 
    # Loop through all languages
    for language in languages:
 
        # Load dataset
        file_name=language+"_dataset.json"
        qa_dataset = EmbeddingQAFinetuneDataset.from_json(file_name)
 
        start_time_assessment=time.time()
 
        # Assess embedding score (in terms of hit rate at k=5)
        score = evaluate(qa_dataset, tokenizer, embed_model, model_spec['normalize'], model_spec['max_length'], model_spec['pooling_type'])
 
        # Get duration of score assessment
        duration_assessment = time.time()-start_time_assessment
 
        results.append([language, model_name, score, duration_assessment])
 
 df_results = pd.DataFrame(results, columns = ["Language" ,"Embedding model", "MRR", "Duration"])