Update utils/func.py

utils/func.py

@@ -3,33 +3,44 @@ import faiss
 import numpy as np
 import datasets
 from transformers import AutoTokenizer, AutoModel
+from config_data.config import Config, load_config
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
 
+config: Config = load_config()
+
+
+def embed_bert_cls(
+        text: str, 
+        model: AutoModel, 
+        tokenizer: AutoTokenizer
+) -> np.ndarray:
 
-def embed_bert_cls(text, model, tokenizer):
     t = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
     with torch.no_grad():
         model_output = model(**{k: v.to(model.device) for k, v in t.items()})
     embeds = model_output.last_hidden_state[:, 0, :]
     embeds = torch.nn.functional.normalize(embeds)
+
     return embeds[0].cpu().numpy()
 
 
-def get_ranked_docs(query, vec_query_base, data,
-    bi_model, bi_tok, cross_model, cross_tok):
+def get_ranked_docs(
+    query: str, vec_query_base: np.ndarray, data: datasets,
+    bi_model: AutoModel, bi_tok: AutoTokenizer, 
+    cross_model: AutoModel, cross_tok: AutoTokenizer
+) -> str:
 
     vec_shape = vec_query_base.shape[1]
     index = faiss.IndexFlatL2(vec_shape)
     index.add(vec_query_base)
     xq = embed_bert_cls(query, bi_model, bi_tok)
-    _, I = index.search(xq.reshape(1, vec_shape), 50)
+    _, I = index.search(xq.reshape(1, vec_shape), 50)  # corpus contains 50 similar queries
     corpus = [data[int(i)]['answer'] for i in I[0]]
 
     queries = [query] * len(corpus)
     tokenized_texts = cross_tok(
         queries, corpus, max_length=128, padding=True, truncation=True, return_tensors="pt"
-    ).to(device)
+    ).to(config.model.device)
 
     with torch.no_grad():
         model_output = cross_model(
@@ -43,7 +54,7 @@ def get_ranked_docs(query, vec_query_base, data,
     return corpus[scores_ix[0][0]]
 
 
-def load_dataset(url='ekaterinatao/house_md_context3'):
+def load_dataset(url: str=config.data.dataset) -> datasets:
 
     dataset = datasets.load_dataset(url, split='train')
     house_dataset = dataset.filter(lambda row: row['labels'] == 0)
@@ -51,7 +62,7 @@ def load_dataset(url='ekaterinatao/house_md_context3'):
     return house_dataset
 
 
-def load_cls_base(url='ekaterinatao/house_md_cls_embeds'):
+def load_cls_base(url: str=config.data.cls_vec) -> np.array:
 
     cls_dataset = datasets.load_dataset(url, split='train')
     cls_base = np.stack([embed['cls_embeds'] for embed in cls_dataset])
@@ -59,7 +70,9 @@ def load_cls_base(url='ekaterinatao/house_md_cls_embeds'):
     return cls_base
 
 
-def load_bi_enc_model(checkpoint='ekaterinatao/house-md-bot-bert-bi-encoder'):
+def load_bi_enc_model(
+        checkpoint: str=config.model.bi_checkpoint
+) -> tuple[AutoTokenizer, AutoModel]:
 
     bi_model = AutoModel.from_pretrained(checkpoint)
     bi_tok = AutoTokenizer.from_pretrained(checkpoint)
@@ -67,24 +80,11 @@ def load_bi_enc_model(checkpoint='ekaterinatao/house-md-bot-bert-bi-encoder'):
     return bi_model, bi_tok
 
 
-def load_cross_enc_model(checkpoint='ekaterinatao/house-md-bot-bert-cross-encoder'):
+def load_cross_enc_model(
+        checkpoint: str=config.model.cross_checkpoint
+) -> tuple[AutoTokenizer, AutoModel]:
 
     cross_model = AutoModel.from_pretrained(checkpoint)
     cross_tok = AutoTokenizer.from_pretrained(checkpoint)
 
-    return cross_model, cross_tok
-
-
-def get_answer(message):
-
-    dataset = load_dataset()
-    cls_base = load_cls_base()
-    bi_enc_model = load_bi_enc_model()
-    cross_enc_model = load_cross_enc_model()
-
-    answer = get_ranked_docs(
-        query=message, vec_query_base=cls_base, data=dataset,
-        bi_model=bi_enc_model[0], bi_tok=bi_enc_model[1],
-        cross_model=cross_enc_model[0], cross_tok=cross_enc_model[1]
-    )
-    return answer
+    return cross_model, cross_tok