from rank_bm25 import BM25Plus

corpus = [
    "Hello there good man!",
    "It is quite windy in London",
    "How is the weather today?"
]

# 원래 Sparse retriever 할 시, TfidfVectorizer에서 tokenize를 넣어 자동으로 적용됐지만 BM25는 tokenize를 따로 해줘야 합니다!
tokenized_corpus = [doc.split(" ") for doc in corpus] 

bm25 = BM25Plus(tokenized_corpus) # <rank_bm25.BM25Plus at 0x7fc5a7a3cd00>

queries = ["windy London", "cold weather","what is weather"]
tokenized_queries= [i.split(" ") for i in queries]

doc_scores = []
doc_indices = []
for idx,i in enumerate(tokenized_queries):
    print("질문: "+ queries[idx])
   
   # get_scores 함수를 통해 passage embedding과 query vector간의 계산 완료!
    scores = bm25.get_scores(i)
    print("Scores: ",scores)

    sorted_score = np.sort(scores)[::-1]
    sorted_indices = np.argsort(scores)[::-1]
    print("오름차순 index: ",sorted_indices)

질문: windy London
Scores:  [2.77258872 5.3162481  2.77258872]
오름차순 index:  [1 2 0]
질문: cold weather
Scores:  [1.38629436 1.38629436 2.77258872]
오름차순 index:  [2 1 0]
질문: what is weather
Scores:  [2.07944154 2.71535639 4.15888308]
오름차순 index:  [2 1 0]

# 기존 init 함수에서 추가한 부분!
self.BM25 = None # BM25 class의 instant를 저장
self.tokenizer = tokenize_fn #BM25 instant 생성 전에, 사용할 tokenizer 저장 

# Pickle을 저장합니다.
pickle_name = f"BM25_embedding.bin"
# BM class instant(BM passage embedding)을 저장하거나 저장된 경로 설정
bm_emd_path = os.path.join(self.data_path, pickle_name)
        
# BM25 존재하면 가져오기
if os.path.isfile(bm_emd_path):
    with open(bm_emd_path, "rb") as file:
         self.BM25 = pickle.load(file)            
    print("BM25_class_instant pickle load.")
        
# https://github.com/dorianbrown/rank_bm25 -> initalizing 부분
# BM25 존재 하지 않으면, tokenizer 한 후, BM25Plus로 passage embedding?
else:
    print("Build BM25_class_instant")
    # BM25는 어떤 text 전처리 X ->  tokenize 따로 한 후에, BM25 클래스의 인스턴스를 생성
    tokenized_contexts= [self.tokenizer(i) for i in self.contexts]
    self.BM25 = BM25Plus(tokenized_contexts)           
    with open(bm_emd_path, "wb") as file:
         pickle.dump(self.BM25, file)
    print("BM25_class_instant pickle saved.")

tokenized_query = self.tokenizer(query) 
        
# ex. array([2.77258872, 5.3162481 , 2.77258872])
# passage embedding과 query vector간의 계산 완료!
doc_scores = self.BM25.get_scores(tokenized_query)
        
# score 높은 순으로 index 정렬
doc_indices=np.argsort(-doc_scores) 
# top_k만큼 뽑기!
return doc_scores[doc_indices[:k]], doc_indices[:k]

print("Build BM25 score, indices")
tokenized_queries= [self.tokenizer(i) for i in queries]        
doc_scores = []
doc_indices = []
for i in tqdm(tokenized_queries):
    scores = self.BM25.get_scores(i)
    ## score가 0이거나 너무 낮은 점수가 꽤 있는 걸 확인하여 score 제한을 어느 정도 걸어서 받아도 될 것 같습니다!
    ## 저는 가장 높은 점수*0.5 기준으로 넘으면 넣는 식으로 했었습니다!
   
    # score 오름차순
    sorted_scores = np.sort(scores)[::-1]
    # index 오름차순
    sorted_indices = np.argsort(scores)[::-1]
    
    doc_scores.append(sorted_scores[:k])
    doc_indices.append(sorted_indices[:k])
return doc_scores, doc_indices

# 바뀐 부분!
## assert self.p_embedding is not None, "get_sparse_embedding() 메소드를 먼저 수행해줘야합니다."
assert self.BM25 is not None, "get_sparse_BM25() 메소드를 먼저 수행해줘야합니다."

if isinstance(query_or_dataset, str):
    # 바뀐 부분!
    ## doc_scores, doc_indices = self.get_relevant_doc(query_or_dataset, k=topk)
    doc_scores, doc_indices = self.get_relevant_doc_BM25(query_or_dataset, k=topk)
    print("[Search query]\n", query_or_dataset, "\n")

    for i in range(topk):
        print(f"Top-{i+1} passage with score {doc_scores[i]:4f}")
        print(self.contexts[doc_indices[i]])

    return (doc_scores, [self.contexts[doc_indices[i]] for i in range(topk)])

elif isinstance(query_or_dataset, Dataset):
    # Retrieve한 Passage를 pd.DataFrame으로 반환합니다.
    total = []
    with timer("query exhaustive search"):
        # 바뀐 부분!
        ## doc_scores, doc_indices = self.get_relevant_doc_bulk(query_or_dataset["question"], k=topk)
        doc_scores, doc_indices = self.get_relevant_doc_bulk_BM25(query_or_dataset['question'], k=topk)
    for idx, example in enumerate(
        tqdm(query_or_dataset, desc="BM25 retrieval: ")
    ):
        
        tmp = {
            # Query와 해당 id를 반환합니다.
            "question": example["question"],
            "id": example["id"],
            # Retrieve한 Passage의 id, context를 반환합니다.
            "context_id": doc_indices[idx],
            "context": " ".join(
                [self.contexts[pid] for pid in doc_indices[idx]]
            ),
        }
        if "context" in example.keys() and "answers" in example.keys():
            # validation 데이터를 사용하면 ground_truth context와 answer도 반환합니다.
            tmp["original_context"] = example["context"]
            tmp["answers"] = example["answers"]
        total.append(tmp)
        
    cqas = pd.DataFrame(total)
    return cqas

if data_args.train_retrieval:
    retriever = SparseRetrieval(tokenize_fn=tokenizer.tokenize,
                                data_path="../data",
                                context_path="wikipedia_documents.json")
    retriever.get_sparse_BM25()

    
# do_train mrc model 혹은 do_eval mrc model
if training_args.do_train or training_args.do_eval:
    run_mrc(data_args, training_args, model_args, datasets, tokenizer, model)

# Query에 맞는 Passage들을 Retrieval 합니다.
# retriever 설정
retriever = SparseRetrieval(
    tokenize_fn=tokenize_fn, data_path=data_path, context_path=context_path
)
    
# Passage Embedding 만들기
#retriever.get_sparse_embedding()
retriever.get_sparse_BM25()
df = retriever.retrieve_BM25(datasets['validation'], topk=data_args.top_k_retrieval)

kor_begin = 44032 # 한글 시작
kor_end = 55203 # 한글 끝
chosung_base = 588 # 초성, base* chosung_list[index]
jungsung_base = 28 # 중성 base* chosung_list[index]
jaum_begin = 12593 #자음 시작
jaum_end = 12622 #자음 끝
moum_begin = 12623 #모음 시작
moum_end = 12643 #모음 끝

chosung_list = [ 'ㄱ', 'ㄲ', 'ㄴ', 'ㄷ', 'ㄸ', 'ㄹ', 'ㅁ', 'ㅂ', 'ㅃ', 
        'ㅅ', 'ㅆ', 'ㅇ' , 'ㅈ', 'ㅉ', 'ㅊ', 'ㅋ', 'ㅌ', 'ㅍ', 'ㅎ']

jungsung_list = ['ㅏ', 'ㅐ', 'ㅑ', 'ㅒ', 'ㅓ', 'ㅔ', 
        'ㅕ', 'ㅖ', 'ㅗ', 'ㅘ', 'ㅙ', 'ㅚ', 
        'ㅛ', 'ㅜ', 'ㅝ', 'ㅞ', 'ㅟ', 'ㅠ', 
        'ㅡ', 'ㅢ', 'ㅣ']

jongsung_list = [
    ' ', 'ㄱ', 'ㄲ', 'ㄳ', 'ㄴ', 'ㄵ', 'ㄶ', 'ㄷ',
        'ㄹ', 'ㄺ', 'ㄻ', 'ㄼ', 'ㄽ', 'ㄾ', 'ㄿ', 'ㅀ', 
        'ㅁ', 'ㅂ', 'ㅄ', 'ㅅ', 'ㅆ', 'ㅇ', 'ㅈ', 'ㅊ', 
        'ㅋ', 'ㅌ', 'ㅍ', 'ㅎ']

jaum_list = ['ㄱ', 'ㄲ', 'ㄳ', 'ㄴ', 'ㄵ', 'ㄶ', 'ㄷ', 'ㄸ', 'ㄹ', 
              'ㄺ', 'ㄻ', 'ㄼ', 'ㄽ', 'ㄾ', 'ㄿ', 'ㅀ', 'ㅁ', 'ㅂ', 
              'ㅃ', 'ㅄ', 'ㅅ', 'ㅆ', 'ㅇ', 'ㅈ', 'ㅉ', 'ㅊ', 'ㅋ', 'ㅌ', 'ㅍ', 'ㅎ']

moum_list = ['ㅏ', 'ㅐ', 'ㅑ', 'ㅒ', 'ㅓ', 'ㅔ', 'ㅕ', 'ㅖ', 'ㅗ', 'ㅘ', 
              'ㅙ', 'ㅚ', 'ㅛ', 'ㅜ', 'ㅝ', 'ㅞ', 'ㅟ', 'ㅠ', 'ㅡ', 'ㅢ', 'ㅣ']

def leven_compose(chosung, jungsung, jongsung):
    char = chr(
        kor_begin +
        chosung_base * chosung_list.index(chosung) +
        jungsung_base * jungsung_list.index(jungsung) +
        jongsung_list.index(jongsung)
    )
    return char

def leven_decompose(c):
    # print('decompose에서 호출')
    # print(c)
    if not leven_character_is_korean(c):
        return c.upper() # return None에서 수정
    i = ord(c)
    if (jaum_begin <= i <= jaum_end):
        return (c, ' ', ' ')
    if (moum_begin <= i <= moum_end):
        return (' ', c, ' ')

    # decomposition rule
    i -= kor_begin
    cho  = i // chosung_base
    jung = ( i - cho * chosung_base ) // jungsung_base 
    jong = ( i - cho * chosung_base - jung * jungsung_base )    
    return (chosung_list[cho], jungsung_list[jung], jongsung_list[jong])

def leven_character_is_korean(c):
    # print('characteriskorean에서 호출')
    # print(c)
    i = ord(c)
    return ((kor_begin <= i <= kor_end) or
            (jaum_begin <= i <= jaum_end) or
            (moum_begin <= i <= moum_end))

def levenshtein(s1, s2, cost=None, debug=False):
    if len(s1) < len(s2):
        return levenshtein(s2, s1, debug=debug)

    if len(s2) == 0:
        return len(s1)

    if cost is None:
        cost = {}

    # changed
    def substitution_cost(c1, c2):
        if c1 == c2:
            return 0
        return cost.get((c1, c2), 1)

    previous_row = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            # Changed
            substitutions = previous_row[j] + substitution_cost(c1, c2)
            current_row.append(min(insertions, deletions, substitutions))

        if debug:
            print(current_row[1:])

        previous_row = current_row

    return previous_row[-1]

def jamo_levenshtein(s1, s2, debug=False):
    if len(s1) < len(s2):
        return jamo_levenshtein(s2, s1, debug)

    if len(s2) == 0:
        return len(s1)

    def substitution_cost(c1, c2):
        if c1 == c2:
            return 0
        return levenshtein(leven_decompose(c1), leven_decompose(c2))/3

    previous_row = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            # Changed
            substitutions = previous_row[j] + substitution_cost(c1, c2)
            current_row.append(min(insertions, deletions, substitutions))

        # if debug:
        # print(['%.3f'%v for v in current_row[1:]])

        previous_row = current_row

    return previous_row[-1]

def return_correct_name_by_leven(data_list, target_name):
  def get_target(data_list,target_name):
    return list(map(lambda x,: jamo_levenshtein(x,target_name), data_list))
  leven_dist = get_target(data_list, target_name)
  return data_list[np.argmin(leven_dist)]

return_correct_name_by_leven(["하이닉스","삼성","엘지"], "ㅎ닉스")

def return_correct_name_by_leven(data_list, target_name):
  def get_target(data_list,target_name):
    return list(map(lambda x,: jamo_levenshtein(x,target_name), data_list))
  leven_dist = get_target(data_list, target_name)
  return data_list[np.argmin(leven_dist)]

def jamo_levenshtein(s1, s2, debug=False):
    if len(s1) < len(s2):
        return jamo_levenshtein(s2, s1, debug)

    if len(s2) == 0:
        return len(s1)

    def substitution_cost(c1, c2):
        if c1 == c2:
            return 0
        return levenshtein(leven_decompose(c1), leven_decompose(c2))/3

    previous_row = range(len(s2) + 1)
    for i, c1 in enumerate(s1):
        current_row = [i + 1]
        for j, c2 in enumerate(s2):
            insertions = previous_row[j + 1] + 1
            deletions = current_row[j] + 1
            # Changed
            substitutions = previous_row[j] + substitution_cost(c1, c2)
            current_row.append(min(insertions, deletions, substitutions))

        # if debug:
        # print(['%.3f'%v for v in current_row[1:]])

        previous_row = current_row

    return previous_row[-1]

def leven_decompose(c):
    if not leven_character_is_korean(c):
        return c.upper() # return None에서 수정
    i = ord(c)
    if (jaum_begin <= i <= jaum_end):
        return (c, ' ', ' ')
    if (moum_begin <= i <= moum_end):
        return (' ', c, ' ')

    # decomposition rule
    i -= kor_begin
    cho  = i // chosung_base
    jung = ( i - cho * chosung_base ) // jungsung_base 
    jong = ( i - cho * chosung_base - jung * jungsung_base )    
    return (chosung_list[cho], jungsung_list[jung], jongsung_list[jong])

def leven_character_is_korean(c):
    # print('characteriskorean에서 호출')
    # print(c)
    i = ord(c)
    return ((kor_begin <= i <= kor_end) or
            (jaum_begin <= i <= jaum_end) or
            (moum_begin <= i <= moum_end))

def leven_compose(chosung, jungsung, jongsung):
    char = chr(
        kor_begin +
        chosung_base * chosung_list.index(chosung) +
        jungsung_base * jungsung_list.index(jungsung) +
        jongsung_list.index(jongsung)
    )
    return char

return_correct_name_by_leven(["하이닉스","삼성","엘지"], "ㅎ닉스")
>>> '하이닉스'

pip install jamo

from transformers import AdamW, get_cosine_with_hard_restarts_schedule_with_warmup
    def create_optimizer_and_scheduler(self, num_training_steps: int, num_cycles:int = 1, another_scheduler_flag=False):
        if not another_scheduler_flag:
            self.create_optimizer()
            self.create_scheduler(num_training_steps=num_training_steps, optimizer=self.optimizer)
        else:
            optimizer_kwargs = {
                    "betas": (self.args.adam_beta1, self.args.adam_beta2),
                    "eps": self.args.adam_epsilon,
                    "lr" : self.args.learning_rate,
                }

            self.optimizer = AdamW(self.model.parameters(), **optimizer_kwargs)
            self.lr_scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(
                                self.optimizer, num_warmup_steps=self.args.warmup_steps, 
                                num_training_steps= num_training_steps,
                                num_cycles = num_cycles)

trainer = QuestionAnsweringTrainer( 
        model=model,
        args=training_args,
        train_dataset=train_dataset if training_args.do_train else None,
        eval_dataset=eval_dataset if training_args.do_eval else None,
        eval_examples=datasets["validation"] if training_args.do_eval else None,
        tokenizer=tokenizer,
        data_collator=data_collator,
        post_process_function=post_processing_function,
        compute_metrics=compute_metrics,
        )
    
   total_steps = math.ceil(len(train_dataset)/training_args.per_device_train_batch_size)
   trainer.create_optimizer_and_scheduler(total_steps, data_args.num_cycles, data_args.another_scheduler_flag)

    another_scheduler_flag :bool = field(
        default=False,
        metadata={"help": "create another scheduler"}
    )
    num_cycles :bool = field(
        default=1,
        metadata={"help": "cycles for get_cosine_schedule_with_warmup"}

#self.lr_scheduler = get_cosine_with_hard_restarts_schedule_with_warmup(
#                            self.optimizer, num_warmup_steps=self.args.warmup_steps, 
#                            num_training_steps= num_training_steps,
#                            num_cycles = num_cycles)

self.lr_scheduler = CosineAnnealingWarmupRestarts(self.optimizer, first_cycle_steps=100, cycle_mult=1.0, max_lr=5e-05, min_lr=1e-07, warmup_steps=50, gamma=0.5)

!pip install 'git+https://github.com/katsura-jp/pytorch-cosine-annealing-with-warmup'