feat: add pooling: cls/mean

app.py

@@ -21,6 +21,7 @@ def read_data(filepath: str) -> Optional[pd.DataFrame]:
 def process(
         task_name: str,
         model_name: str,
+        pooling: str,
         text: str,
         file=None,
 ) -> (None, pd.DataFrame, str):
@@ -37,9 +38,9 @@ def process(
 
     # process
     if task_name == 'Originality':
-        df = pipeline.p0_originality(df, model_name)
+        df = pipeline.p0_originality(df, model_name, pooling)
     elif task_name == 'Flexibility':
-        df = pipeline.p1_flexibility(df, model_name)
+        df = pipeline.p1_flexibility(df, model_name, pooling)
     else:
         raise Exception('Task not supported')
 
@@ -62,6 +63,11 @@ model_name_dropdown = gr.components.Dropdown(
     value=list_models[0],
     choices=list_models
 )
+pooling_dropdown = gr.components.Dropdown(
+    label='Pooling',
+    value='mean',
+    choices=['mean', 'cls']
+)
 text_input = gr.components.Textbox(
     value=open('data/example_xlm.csv', 'r').read(),
     lines=10,
@@ -75,7 +81,7 @@ file_output = gr.components.File(label='Output File', file_types=['.csv', '.xlsx
 
 app = gr.Interface(
     fn=process,
-    inputs=[task_name_dropdown, model_name_dropdown, text_input, file_input],
+    inputs=[task_name_dropdown, model_name_dropdown, pooling_dropdown, text_input, file_input],
     outputs=[text_output, dataframe_output, file_output],
     description=open('data/description.txt', 'r').read()
 )

utils/models.py

@@ -2,6 +2,7 @@ from functools import lru_cache
 
 import torch
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
 
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
@@ -10,7 +11,9 @@ list_models = [
     'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2',
     'sentence-transformers/all-mpnet-base-v2',
     'sentence-transformers/all-MiniLM-L12-v2',
-    'cyclone/simcse-chinese-roberta-wwm-ext'
+    'cyclone/simcse-chinese-roberta-wwm-ext',
+    'bert-base-chinese',
+    'IDEA-CCNL/Erlangshen-SimCSE-110M-Chinese',
 ]
 
 
@@ -18,8 +21,60 @@ class SBert:
     def __init__(self, path):
         print(f'Loading model from {path} ...')
         self.model = SentenceTransformer(path, device=DEVICE)
+        # from pprint import pprint
+        # pprint(self.model.__dict__)
 
     @lru_cache(maxsize=10000)
     def __call__(self, x) -> torch.Tensor:
         y = self.model.encode(x, convert_to_tensor=True)
         return y
+
+
+class ModelWithPooling:
+    def __init__(self, path):
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModel.from_pretrained(path)
+
+    @lru_cache(maxsize=10000)
+    @torch.no_grad()
+    def __call__(self, text: str, pooling='mean'):
+        inputs = self.tokenizer(text, padding=True, truncation=True, return_tensors="pt")
+        outputs = self.model(**inputs, output_hidden_states=True)
+
+        if pooling == 'cls':
+            o = outputs.last_hidden_state[:, 0]  # [b, h]
+
+        elif pooling == 'pooler':
+            o = outputs.pooler_output  # [b, h]
+
+        elif pooling in ['mean', 'last-avg']:
+            last = outputs.last_hidden_state.transpose(1, 2)  # [b, h, s]
+            o = torch.avg_pool1d(last, kernel_size=last.shape[-1]).squeeze(-1)  # [b, h]
+
+        elif pooling == 'first-last-avg':
+            first = outputs.hidden_states[1].transpose(1, 2)  # [b, h, s]
+            last = outputs.hidden_states[-1].transpose(1, 2)  # [b, h, s]
+            first_avg = torch.avg_pool1d(first, kernel_size=last.shape[-1]).squeeze(-1)  # [b, h]
+            last_avg = torch.avg_pool1d(last, kernel_size=last.shape[-1]).squeeze(-1)  # [b, h]
+            avg = torch.cat((first_avg.unsqueeze(1), last_avg.unsqueeze(1)), dim=1)  # [b, 2, h]
+            o = torch.avg_pool1d(avg.transpose(1, 2), kernel_size=2).squeeze(-1)  # [b, h]
+
+        else:
+            raise Exception(f'Unknown pooling {pooling}')
+
+        o = o.squeeze(0)
+        return o
+
+
+def test_sbert():
+    m = SBert('bert-base-chinese')
+    o = m('hello')
+    print(o.size())
+    assert o.size() == (768,)
+
+
+def test_hf_model():
+    m = ModelWithPooling('IDEA-CCNL/Erlangshen-SimCSE-110M-Chinese')
+    o = m('hello', pooling='cls')
+    print(o.size())
+    assert o.size() == (768,)

utils/pipeline.py

@@ -3,10 +3,10 @@ from typing import List
 import pandas as pd
 from sentence_transformers.util import cos_sim
 
-from utils.models import SBert
+from utils.models import ModelWithPooling
 
 
-def p0_originality(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
+def p0_originality(df: pd.DataFrame, model_name: str, pooling: str) -> pd.DataFrame:
     """
     row-wise
     :param df:
@@ -15,11 +15,11 @@ def p0_originality(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
     """
     assert 'prompt' in df.columns
     assert 'response' in df.columns
-    model = SBert(model_name)
+    model = ModelWithPooling(model_name)
 
     def get_cos_sim(prompt: str, response: str) -> float:
-        prompt_vec = model(prompt)
-        response_vec = model(response)
+        prompt_vec = model(text=prompt, pooling=pooling)
+        response_vec = model(text=response, pooling=pooling)
         score = cos_sim(prompt_vec, response_vec).item()
         return score
 
@@ -27,7 +27,7 @@ def p0_originality(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
     return df
 
 
-def p1_flexibility(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
+def p1_flexibility(df: pd.DataFrame, model_name: str, pooling: str) -> pd.DataFrame:
     """
     group-wise
     :param df:
@@ -37,10 +37,10 @@ def p1_flexibility(df: pd.DataFrame, model_name: str) -> pd.DataFrame:
     assert 'prompt' in df.columns
     assert 'response' in df.columns
     assert 'id' in df.columns
-    model = SBert(model_name)
+    model = ModelWithPooling(model_name)
 
     def get_flexibility(responses: List[str]) -> float:
-        responses_vec = [model(_) for _ in responses]
+        responses_vec = [model(text=_, pooling=pooling) for _ in responses]
         score = 0
         for i in range(len(responses_vec) - 1):
             score += 1 - cos_sim(responses_vec[i], responses_vec[i + 1]).item()