embeddings based project match, other optims

app.py

@@ -4,6 +4,7 @@ import re
 from types import SimpleNamespace
 
 import gradio as gr
+from itertools import chain
 import pandas as pd
 import spaces
 import spacy
@@ -18,6 +19,7 @@ from transformers import (
     BertTokenizer,
     set_seed,
 )
+from sentence_transformers import SentenceTransformer, util
 
 from methods import gdc_api_calls, utilities
 
@@ -74,6 +76,10 @@ intent_tok = AutoTokenizer.from_pretrained(
 intent_model = BertForSequenceClassification.from_pretrained(model_id, token=hf_TOKEN)
 intent_model = intent_model.to("cuda").eval()
 
+# load sentence transformer model to test cancer embeddings
+st_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+st_model = st_model.to("cuda")
+
 
 print("loading gdc genes and mutations")
 gdc_genes_mutations = utilities.load_gdc_genes_mutations_hf(hf_TOKEN)
@@ -112,6 +118,17 @@ def infer_gene_entities_from_query(query):
 
 
 
+def get_project_embeddings():
+    project_rows = []
+    for k,v in project_mappings.items():
+        new_v = [item.replace(',', '') for item in v]
+        combined = ','.join([k] + new_v)
+        project_rows.append(combined)
+    row_embeddings = model.encode(project_rows, convert_to_tensor=True)
+    return project_rows, row_embeddings
+
+
+
 def check_if_project_id_in_query(query):
     # check if mention of project keys
     # e.g. TCGA-BRCA in query
@@ -147,11 +164,31 @@ def proj_id_and_partial_match(query, initial_cancer_entities):
 
 
 
+def get_top_k_cancer_entities(query, row_embeddings, project_rows, top_k=20):
+    top_cancer_entities = []
+    query_embedding = st_model.encode(query, convert_to_tensor=True)
+    cosine_scores = util.cos_sim(query_embedding, row_embeddings)
+    top_results = torch.topk(cosine_scores, k=top_k)
+    top_results_indices = top_results.indices.tolist()
+    top_results_scores = top_results.values.tolist()
+    print(top_results_scores)
+    for idx, score in enumerate(top_results_scores[0]):
+        if score > 0.5:
+            row_idx = top_results_indices[0][idx]
+            print('best row, score: {} {}'.format(project_rows[row_idx], score))
+            top_cancer_entities.append([project_rows[row_idx], score])
+    try:
+        top_projects = [sublist[0].split(',')[0] for sublist in top_cancer_entities]
+    except Exception as e:
+        top_projects = []
+    return top_projects
+
+
 
 @utilities.timeit
 def postprocess_cancer_entities(initial_cancer_entities, query):
     # print('initial cancer entities {}'.format(initial_cancer_entities))
-    # print('check if GDC project-id mentioned in query')
+    project_rows, row_embeddings = get_project_embeddings()
     final_entities = check_if_project_id_in_query(query)
     if final_entities:
         return final_entities
@@ -171,6 +208,13 @@ def postprocess_cancer_entities(initial_cancer_entities, query):
                 final_entities = proj_id_and_partial_match(
                     query, initial_cancer_entities
                 )
+            # try embedding based match
+            if not final_entities:
+                print('Test embedding based match')
+                for i in initial_cancer_entities:
+                    c_entities = get_top_k_cancer_entities(i, row_embeddings, project_rows)
+                    final_entities.append(c_entities)
+                final_entities = list(chain.from_iterable(final_entities))
         else:
             # no initial_cancer_entities
             # check project_mappings keys/values for matches with query terms
@@ -400,28 +444,20 @@ def batch_test(query):
 def get_prefinal_response(row):
     try:
         query = row["questions"]
-        genes = ','.join(row['gene_entities'])
         gdc_result = row["gdc_result"]
     except Exception as e:
         print(f"unable to retrieve query: {query} or gdc_result: {gdc_result}")
 
-    intent = intent_expansion[row['intent']]
-    
-    print("\nStep 6: Construct LLM prompts for llama-3B\n")
-    descriptive_prompt = construct_modified_query_description(genes, intent)
+    print("\nStep 6: Construct LLM prompts (percentage) for llama-3B\n")
     percentage_prompt = construct_modified_query_percentage(query, gdc_result)
     
-    print("\nStep 7: Generate LLM response R on query augmented prompts\n")
-    descriptive_response = generate_descriptive_response(descriptive_prompt)
-    if not descriptive_response.endswith('.'):
-        descriptive_response += '.'
+    print("\nStep 7: Generate LLM response R (percentage) on query augmented prompts\n")
 
     percentage_response = generate_percentage_response(percentage_prompt)
     percentage_response = re.sub(
         r'final response', 'frequency for your query', percentage_response)
     return pd.Series([
-        descriptive_prompt, percentage_prompt, 
-        descriptive_response, percentage_response
+        percentage_prompt, percentage_response
         ])
 
 
@@ -630,7 +666,21 @@ def execute_pipeline(question: str):
             ]
         ] = df["questions"].apply(lambda x: batch_test(x))
         df_exploded = df.explode("gdc_result", ignore_index=True)
-        df_exploded[["descriptive_prompt", "percentage_prompt", "descriptive_response", "percentage_response"]] = df_exploded.apply(
+
+        # generate descriptive response once based on genes and intent
+        print("\nStep 6: Construct LLM prompts (descriptive) for llama-3B\n")
+        intent = intent_expansion[df['intent'].iloc[0]]
+        genes = ','.join(df['gene_entities'].iloc[0])
+        descriptive_prompt = construct_modified_query_description(genes, intent)
+
+        print("\nStep 7: Generate LLM response R (descriptive) on query augmented prompts\n")
+        descriptive_response = generate_descriptive_response(descriptive_prompt, model, tok)
+        if not descriptive_response.endswith('.'):
+            descriptive_response += '.'
+
+        df_exploded[['descriptive_prompt', 'descriptive_response']] = descriptive_prompt, descriptive_response
+
+        df_exploded[["percentage_prompt", "percentage_response"]] = df_exploded.apply(
             lambda x: get_prefinal_response(x), axis=1)