Upload app.py

app.py

@@ -12,6 +12,19 @@ from model import RNN_model
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+# 导入预训练模型和分词器
+model_name = "microsoft/phi-2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# 设置填充令牌，如果分词器没有默认的填充令牌
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModel.from_pretrained(model_name)
+
 # Import data
 df= pd.read_csv('Symptom2Disease.csv')
 df.drop('Unnamed: 0', axis= 1, inplace= True)
@@ -47,17 +60,36 @@ class_names= {0: 'Acne',
               23: 'urinary tract infection'
               }
 
-vectorizer= nltk_u.vectorizer()
-vectorizer.fit(train_data.text)
-
-
-
-# Model and transforms preparation
-model= RNN_model()
-# Load state dict
-model.load_state_dict(torch.load(
-    f= 'pretrained_symtom_to_disease_model.pth',
-    map_location= torch.device('cpu')))
+# 数据预处理
+def preprocess(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    return inputs
+
+# 模型预测逻辑
+def get_prediction(inputs):
+    model.eval()
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.last_hidden_state[:, 0, :]  # 取CLS标记的输出进行分类
+        pred_prob = torch.softmax(logits, dim=1)
+        pred = torch.argmax(pred_prob, dim=1).item()
+        if pred in class_names:
+            return class_names[pred]
+        else:
+            print(f"Warning: Prediction index {pred} not found in class_names.")
+            return "Unknown"  # 或者其他默认的响应
+
+# vectorizer= nltk_u.vectorizer()
+# vectorizer.fit(train_data.text)
+
+
+
+# # Model and transforms preparation
+# model= RNN_model()
+# # Load state dict
+# model.load_state_dict(torch.load(
+#     f= 'pretrained_symtom_to_disease_model.pth',
+#     map_location= torch.device('cpu')))
 # Disease Advice
 disease_advice = {
     'Acne': "Maintain a proper skincare routine, avoid excessive touching of the affected areas, and consider using over-the-counter topical treatments. If severe, consult a dermatologist.",
@@ -90,7 +122,6 @@ howto= """Welcome to the <b>Medical Chatbot</b>, powered by Gradio.
 Currently, the chatbot can WELCOME YOU, PREDICT DISEASE based on your symptoms and SUGGEST POSSIBLE SOLUTIONS AND RECOMENDATIONS, and BID YOU FAREWELL.
 <b>How to Start:</b> Simply type your messages in the textbox to chat with the Chatbot and press enter!<br><br>
 The bot will respond based on the best possible answers to your messages.
-
 """
 
 
@@ -175,17 +206,19 @@ with gr.Blocks(css = """#col_container { margin-left: auto; margin-right: auto;}
     elif message.lower() in goodbyes:
       bot_message= random.choice(goodbye_replies)
     else: 
+        inputs = preprocess(message)
+        bot_message = f"Based on your symptoms, I believe you may have {get_prediction(inputs)}."
       #bot_message= random.choice(goodbye_replies)
-  
-      transform_text= vectorizer.transform([message])
-      transform_text= torch.tensor(transform_text.toarray()).to(torch.float32)
-      model.eval()
-      with torch.inference_mode():
-        y_logits=model(transform_text)
-        pred_prob= torch.argmax(torch.softmax(y_logits, dim=1), dim=1)
+      
+      # transform_text= vectorizer.transform([message])
+      # transform_text= torch.tensor(transform_text.toarray()).to(torch.float32)
+      # model.eval()
+      # with torch.inference_mode():
+      #   y_logits=model(transform_text)
+      #   pred_prob= torch.argmax(torch.softmax(y_logits, dim=1), dim=1)
    
-      test_pred= class_names[pred_prob.item()] 
-      bot_message = f' Based on your symptoms, I believe you are having {test_pred} and I would advice you {disease_advice[test_pred]}'
+      # test_pred= class_names[pred_prob.item()] 
+      # bot_message = f' Based on your symptoms, I believe you are having {test_pred} and I would advice you {disease_advice[test_pred]}'
     chat_history.append((message, bot_message))
     time.sleep(2)
     return "", chat_history
@@ -194,5 +227,4 @@ with gr.Blocks(css = """#col_container { margin-left: auto; margin-right: auto;}
   
 
 # Launch the demo
-demo.launch()
-
+demo.launch()