Update app.py

app.py

@@ -2,84 +2,66 @@
 import pandas as pd
 import numpy as np
 import torch
+import json
+import re
 from sentence_transformers.util import cos_sim
 from sentence_transformers import SentenceTransformer
+from sklearn.feature_extraction.text import CountVectorizer
 import gradio as gr
+import nltk
+nltk.download('stopwords')
+from nltk.corpus import stopwords
 #%%
-# etalon = pd.read_csv("etalon_prod.csv")
-df = pd.read_csv("preprocessed_complaints.csv")
 
 model = SentenceTransformer('sentence-transformers/multi-qa-distilbert-cos-v1')
+russian_stopwords = stopwords.words('russian') + ['ВАШ']
 
-
-unique_complaints = df['Жалобы'].unique()
+with open("data/symptom_checker/top_150_symps_by_spec.json", 'r') as f:
+    symps = json.load(f)
 
 with open("embeddings.npy", 'rb') as f:
-    embeddings = np.load(f)
+    embs = np.load(f)
 
-#%%
-def get_recommend(user_input, 
-                  top_k_spec = 3,
-                  top_k_services = 5,
-                  treshold = 0.8):
-    
-    cols_for_top_k = ["Специальность врача",
-                      "Рекомендуемые специалисты"]
-    
-    usr_embeddings = model.encode(user_input)
+def remove_numbers(text):
+    text = re.sub(r'\d+', '', text)
+    text = re.sub(r'[^\w\s]', '', text)
+    return text.strip()
 
-    cos_similarity = cos_sim(usr_embeddings, embeddings).detach().numpy()
-    sorted_idx = cos_similarity[0].argsort()[::-1]
-    cos_similarity.sort()
-    
-    cos_similarity = cos_similarity[0][::-1]
-    
-    sorted_df = df.loc[sorted_idx].copy()
-    sorted_df['cos_sim'] = cos_similarity
-    sorted_df = sorted_df[sorted_df['cos_sim'] > treshold]
-    
-    result = {}
-    for col in cols_for_top_k:
-        result[col] = sorted_df[col].value_counts()[:top_k_spec].index.tolist()
-    result['Жалобы'] = sorted_df['Жалобы'].value_counts()[:top_k_services].index.tolist()
-    
-    top_k_mkb = sorted_df['Диагноз МКБ'].value_counts()[:top_k_services].index.tolist()
-    result['Диагноз МКБ'] = top_k_mkb
-    
-    categories = ['Инструментальная диагностика', 'Лабораторная диагностика']
+
+
+vectorizer = CountVectorizer(ngram_range=(1, 3),
+                             stop_words=russian_stopwords,
+                             preprocessor=remove_numbers,
+                             )
     
-    top_k_services_lst_by_mkb = []
-    for mkb in top_k_mkb:
-        temp_lst = []
-        slice_df = sorted_df[sorted_df['Диагноз МКБ'] == mkb]
-        for category in categories:
-            top_k_services_in_cat_mkb = slice_df[slice_df['service_name_category'] == category]['Рекомендации по обследованию'].value_counts()[:top_k_services].index.tolist()
-            temp_lst.append({category:top_k_services_in_cat_mkb})
-            
-        top_k_services_lst_by_mkb.append({mkb:temp_lst})
+def get_symptomps_v2(text, treshold = 0.7):
+    try:
+        if isinstance(text, str):
+            text = [text]
             
-    top_k_services_lst = []    
-    
-    for category in categories:
-        slice_df = sorted_df[sorted_df['service_name_category'] == category]
-        list_top_k_services = slice_df['Рекомендации по обследованию'].value_counts()[:top_k_services].index.tolist()
-        top_k_services_lst.append({category:list_top_k_services})
+        X = vectorizer.fit_transform(text)
+        # print(vectorizer.get_feature_names_out())
+        text_emb = model.encode(vectorizer.get_feature_names_out(), batch_size=64)
+        cos_sim_m = util.cos_sim(text_emb, embs).numpy()
+        cos_sim_m = np.where(cos_sim_m > treshold, cos_sim_m, -1)
         
+        arg_max_idx = np.argmax(cos_sim_m, axis=1)
         
-    
-    result['Рекомендации по обследованию'] = top_k_services_lst
-    result['Рекомендации по обследованию по МКБ'] = top_k_services_lst_by_mkb
-
-    return result
+        outputs = []
+        for idx, cos_sim_row in zip(arg_max_idx, cos_sim_m):
+            if cos_sim_row[idx] > 0:
+                outputs.append(symps[idx])
+        if len(outputs) == 0:
+            return ['Симптомы не определены']
+        return np.unique(outputs).tolist()
+    except:
+        return ['Симптомы не определены']
 #%%
 gradio_app = gr.Interface(
     get_recommend,
     inputs=['text',
-            gr.Slider(minimum=1, maximum=10, step=1, label="Топ N специалистов", value=3),
-            gr.Slider(minimum=1, maximum=10, step=1, label="Топ N услуг", value=5),
             gr.Slider(minimum=0, maximum=1, step=0.05, label="Порог релевантности", value=0.8)],
-    outputs=[gr.JSON(label='Рекомендации: ')],
-    # title="Предсказание топ-10 наиболее схожих услуг",
+    outputs=[gr.JSON(label='Симптомы: ')],
     description="Введите услугу:"
 )