update

app.py

@@ -1,6 +1,8 @@
+import torch
 import gradio as gr
 import pandas as pd
-from utils import save_dataframe_to_file
+from utils import save_dataframe_to_file, tokenize_Df
+from model import load_model
 
 with open("./description.md", "r", encoding="utf-8") as file:
     description_text = file.read()
@@ -27,17 +29,26 @@ def process_data(task_name, model_name, pooling_method, input_text=None, file=No
             output += " One-phase Fine-tuned BERT model does not support Appropriateness task."
         else:
             # 读取文件
-            if file.name.endswith('.csv'):
-                df = pd.read_csv(file)
-            else:
-                df = pd.read_excel(file)
+            df = pd.read_csv(file) if file.name.endswith('.csv') else pd.read_excel(file)
             # 检查第一行是否为 "prompt" 和 "response"
             if list(df.columns) == ['prompt', 'response']:
                 dataframe_output = df
             else:
                 df_values = [list(df.columns)] + df.values.tolist()
                 dataframe_output = pd.DataFrame(df_values, columns=['prompt', 'response'])
+
+            # model 运行
+            loaded_net = load_model(model_name, pooling_method)
+            example = tokenize_Df(dataframe_output)
+            with torch.no_grad():  
+                score = loaded_net(example)
+
+            if model_name == "One-phase Fine-tuned BERT":
+                dataframe_output['evaluation'] = score.numpy()
+            else: 
+                dataframe_output['evaluation'] = score[0].numpy() if task_name=='Creativity' else score[1].numpy()
             file_output = save_dataframe_to_file(dataframe_output, file_format="csv") 
+            output += f" Processed {len(dataframe_output)} rows from uploaded file using task: {task_name}, model: {model_name}, pooling: {pooling_method}."
 
     # 情况 3: 只有 file
     elif file is not None:
@@ -45,26 +56,35 @@ def process_data(task_name, model_name, pooling_method, input_text=None, file=No
         if not (file.name.endswith('.csv') or file.name.endswith('.xlsx')):
             output = "File format must be xlsx or csv."
         elif task_name == "Appropriateness" and model_name == "One-phase Fine-tuned BERT":
-            output += " One-phase Fine-tuned BERT model does not support Appropriateness task."            
+            output = " One-phase Fine-tuned BERT model does not support Appropriateness task."            
         else:
             # 读取文件
-            if file.name.endswith('.csv'):
-                df = pd.read_csv(file)
-            else:
-                df = pd.read_excel(file)
+            df = pd.read_csv(file) if file.name.endswith('.csv') else pd.read_excel(file)
+
             # 检查第一行是否为 "prompt" 和 "response"
             if list(df.columns) == ['prompt', 'response']:
                 dataframe_output = df
             else:
                 df_values = [list(df.columns)] + df.values.tolist()
                 dataframe_output = pd.DataFrame(df_values, columns=['prompt', 'response'])
+            
+            # model 运行
+            loaded_net = load_model(model_name, pooling_method)
+            example = tokenize_Df(dataframe_output)
+            with torch.no_grad():
+                score = loaded_net(example)
+            
+            if model_name == "One-phase Fine-tuned BERT":
+                dataframe_output['evaluation'] = score.numpy()
+            else: 
+                dataframe_output['evaluation'] = score[0].numpy() if task_name=='Creativity' else score[1].numpy()
             file_output = save_dataframe_to_file(dataframe_output, file_format="csv") 
             output = f"Processed {len(dataframe_output)} rows from uploaded file using task: {task_name}, model: {model_name}, pooling: {pooling_method}."
 
     # 情况 4: 只有 input_text
     elif input_text is not None:
         if task_name == "Appropriateness" and model_name == "One-phase Fine-tuned BERT":
-            output = " One-phase Fine-tuned BERT model does not support Appropriateness task."
+            output = "One-phase Fine-tuned BERT model does not support Appropriateness task."
         else:
             lines = input_text.strip().split("\n")
             rows = []
@@ -78,10 +98,18 @@ def process_data(task_name, model_name, pooling_method, input_text=None, file=No
                     break
             
             if output == "":
-                if rows[0] == ['prompt', 'response']:
-                    dataframe_output = pd.DataFrame(rows[1:], columns=['prompt', 'response'])
-                else:
-                    dataframe_output = pd.DataFrame(rows, columns=['prompt', 'response'])
+                dataframe_output = pd.DataFrame(rows[1:], columns=['prompt', 'response']) if rows[0] == ['prompt', 'response'] else pd.DataFrame(rows, columns=['prompt', 'response'])
+                
+                # model 运行
+                loaded_net = load_model(model_name, pooling_method)
+                example = tokenize_Df(dataframe_output)
+                with torch.no_grad(): 
+                    score = loaded_net(example)
+
+                if model_name == "One-phase Fine-tuned BERT":
+                    dataframe_output['evaluation'] = score.numpy()
+                else: 
+                    dataframe_output['evaluation'] = score[0].numpy() if task_name=='Creativity' else score[1].numpy()
                 file_output = save_dataframe_to_file(dataframe_output, file_format="csv") 
                 output = f"Processed {len(dataframe_output)} rows of text using task: {task_name}, model: {model_name}, pooling: {pooling_method}."
 

model.py

@@ -1,5 +1,12 @@
+import os
 import torch
 from torch import nn
+from transformers import AutoModel
+from huggingface_hub import hf_hub_download
+
+token=os.getenv("HF_TOKEN")
+repo_id = "Siyunb323/CreativityEvaluation"
+model = AutoModel.from_pretrained("cl-tohoku/bert-base-japanese")
 
 class BERTregressor(nn.Module):
     def __init__(self, bert, hidden_size=768, num_linear=1, dropout=0.1, 
@@ -134,4 +141,17 @@ class BERT2Phase(nn.Module):
             e_pred = self.effectiveness(encoded_X.pooler_output)
             c_pred = self.creativity(encoded_X.pooler_output)
         
-        return (c_pred, e_pred) if not return_attention else (c_pred, e_pred, encoded_X.attentions)
+        return (c_pred, e_pred) if not return_attention else (c_pred, e_pred, encoded_X.attentions)
+    
+def load_model(model_name, pooling_method):
+    pooling = pooling_method if pooling_method == 'cls' else 'pooler'
+    if model_name == "One-phase Fine-tuned BERT":
+        loaded_net = BERTregressor(model, hidden_size=768, num_linear=1, dropout=0.1, o_type=pooling, t_type='C', use_sigmoid=True)
+        filename = 'model' + f"/OnePhase_BERT_{pooling_method}.pth" 
+    elif model_name == "Two-phase Fine-tuned BERT":
+        loaded_net = BERT2Phase(model, hidden_size=768, num_linear=1, dropout=0.1, type=pooling, use_sigmoid=True)
+        filename = 'model' + f"/TwoPhase_BERT_{pooling_method}.pth"
+    model_path = hf_hub_download(repo_id=repo_id, filename=filename, use_auth_token=token)
+    loaded_net.load_state_dict(torch.load(model_path))
+    loaded_net.eval()
+    return loaded_net

utils.py

@@ -1,5 +1,9 @@
 import os
 import tempfile
+import fugashi
+import unidic_lite
+from transformers import AutoTokenizer
+tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese")
 
 def save_dataframe_to_file(dataframe, file_format="csv"):
     temp_dir = tempfile.gettempdir()  # 获取系统临时目录
@@ -8,4 +12,11 @@ def save_dataframe_to_file(dataframe, file_format="csv"):
         dataframe.to_csv(file_path, index=False)
     elif file_format == "xlsx":
         dataframe.to_excel(file_path, index=False)
-    return file_path
+    return file_path
+
+def tokenize_Df(examples):
+    return tokenizer(list(examples['prompt']), list(examples['response']),
+                    return_tensors="pt",
+                    padding='max_length',
+                    max_length=60,
+                    truncation='longest_first')