Update app_phi2.py

app_phi2.py

@@ -12,6 +12,14 @@ 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/phi2-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
+
 # Import data
 df= pd.read_csv('Symptom2Disease.csv')
 df.drop('Unnamed: 0', axis= 1, inplace= True)
@@ -47,17 +55,32 @@ class_names= {0: 'Acne',
               23: 'urinary tract infection'
               }
 
-vectorizer= nltk_u.vectorizer()
-vectorizer.fit(train_data.text)
+# 数据预处理
+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)
+        return class_names[pred.item()]
 
+# 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')))
+
+
+# # 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.",
@@ -174,17 +197,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