Upload folder using huggingface_hub

adapter_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:959ff15616d1fb20a339764f8402ee0837c745a64a2441e173bc258ca6d1da32
+oid sha256:b9257294989d1901fc44650804efd5878092c143986bbb9a12120de62e4773bb
 size 209736952

poly_encoder.py

@@ -0,0 +1,90 @@
+import random
+import numpy as np
+import torch
+from transformers import BertTokenizer, BertModel
+import torch.nn.functional as F
+import nltk
+nltk.download('punkt')
+from nltk.tokenize import sent_tokenize
+
+def set_seed(seed):
+    torch.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+class PolyEncoder(torch.nn.Module):
+    def __init__(self, bert_model_name='klue/bert-base', poly_m=16):
+        super(PolyEncoder, self).__init__()
+        self.poly_m = poly_m
+        self.bert_model = BertModel.from_pretrained(bert_model_name)
+        self.poly_code_embeddings = torch.nn.Embedding(poly_m, self.bert_model.config.hidden_size)
+        
+    def forward(self, context_input_ids, context_attention_mask, question_input_ids, question_attention_mask):
+        # Encode the question
+        question_outputs = self.bert_model(input_ids=question_input_ids, attention_mask=question_attention_mask)
+        question_cls_embeddings = question_outputs.last_hidden_state[:, 0, :]  # CLS token
+
+        # Encode the context
+        context_outputs = self.bert_model(input_ids=context_input_ids, attention_mask=context_attention_mask)
+        context_hidden_states = context_outputs.last_hidden_state
+
+        # Poly codes
+        poly_codes = self.poly_code_embeddings.weight.unsqueeze(0).expand(context_hidden_states.size(0), -1, -1)
+
+        # Context and poly code interactions
+        attention_weights = F.softmax(torch.einsum('bmd,bnd->bmn', context_hidden_states, poly_codes), dim=-1)
+        poly_context_embeddings = torch.einsum('bmn,bmd->bnd', attention_weights, context_hidden_states)
+
+        # Question and poly context interactions
+        scores = torch.einsum('bnd,bmd->bnm', poly_context_embeddings, question_cls_embeddings.unsqueeze(1).expand(-1, self.poly_m, -1))
+
+        # Aggregate scores over poly_m dimension
+        scores = scores.max(dim=1).values
+
+        return scores
+
+def get_top_n_relevant_sentences(context, question, tokenizer, model, top_n):
+    context_sentences = sent_tokenize(context)  # NLTK를 사용하여 문장 분할
+
+    context_inputs = tokenizer(context_sentences, padding=True, truncation=True, return_tensors='pt')
+    question_inputs = tokenizer(question, return_tensors='pt')
+
+    with torch.no_grad():
+        scores = model(context_inputs['input_ids'], context_inputs['attention_mask'], 
+                       question_inputs['input_ids'].expand(len(context_sentences), -1),
+                       question_inputs['attention_mask'].expand(len(context_sentences), -1))
+
+    score_rows, score_cols = scores.shape
+
+    scores_index = scores[:, 0].tolist()
+    indexed_dict = {idx: value for idx, value in enumerate(scores_index)}
+    sorted_dict = dict(sorted(indexed_dict.items(), key=lambda item: item[1], reverse=True))
+    sorted_data = sorted(sorted_dict.items(), key=lambda item: item[1], reverse=True)
+    top_n_keys = list(sorted_dict.keys())[:top_n]
+    unique_values = set()
+    top_keys = []
+
+    for key, value in sorted_data:
+        if value not in unique_values:
+            unique_values.add(value)
+            top_keys.append(key)
+        if len(top_keys) == top_n:
+            break
+
+    top_n_sentences = [context_sentences[idx] for idx in top_keys]
+    return top_n_sentences
+
+# 예제 실행 함수
+def run_example(context, question):
+    # 모델 및 토크나이저 로드를 전역 변수로 설정
+    tokenizer = BertTokenizer.from_pretrained('klue/bert-base')
+    model = PolyEncoder(bert_model_name='klue/bert-base')
+
+    top_n_sentences = get_top_n_relevant_sentences(context, question, tokenizer, model, top_n=5)
+    sentences = ""
+    for sentence in top_n_sentences:
+        sentences+=sentence
+    print(sentences)
+    return sentences