from ragdb import TextRAG

from langchain_huggingface import HuggingFaceEmbeddings
embed_model = HuggingFaceEmbeddings(
        model_name = 'alibaba-nlp/gte-multilingual-base',
        model_kwargs = {'device': 'cuda', 'trust_remote_code': True},
        encode_kwargs = {'normalize_embeddings': False}
    )

vectorstore_path = 'rag_index_md'
rag = TextRAG(embed_model=embed_model,
              vectorstore_dir=vectorstore_path)

from evaluator import Evaluator
from chatbot import Chatbot

qa_dir = r"C:\Users\vuvan\Desktop\An_Plaza\ViMedLLM\Vietnamese-Medical-LLM\dataset\QA Data\MedAB\all_data.jsonl"
evaluator = Evaluator(qa_dir = qa_dir,
                      chatbot = Chatbot(model_name="mistral", max_token=1000),
                      rag = rag,
                      search_type = "similarity",
                      log = True)

import json
with open(qa_dir, 'r', encoding="utf-8") as file:
    data = [json.loads(line) for line in file]
data = data[:2]
questions = [item['question'] for item in data]
answers = [item['answer'] for item in data]
choices = [[item['A'], item['B'], item['C'], item['D'], item['E']] for item in data]
print("start:")
acc = evaluator.eval(questions, choices, answers, max_workers=3, suppress_error=True, k=4, threshold=0.5)
print(acc)