using chinese BERT if crosswoz is used

convlab/policy/vector/vector_base.py

@@ -28,7 +28,7 @@ class VectorBase(Vector):
         self.ontology = load_ontology(dataset_name)
         try:
             # execute to make sure that the database exists or is downloaded otherwise
-            if dataset_name == "multiwoz21":
+            if dataset_name == "multiwoz21" or dataset_name == "crosswoz":
                 load_database(dataset_name)
             # the following two lines are needed for pickling correctly during multi-processing
             exec(f'from data.unified_datasets.{dataset_name}.database import Database')

convlab/policy/vtrace_DPT/transformer_model/node_embedder.py

@@ -2,7 +2,7 @@ import os, json, logging
 import torch
 import torch.nn as nn
 
-from transformers import RobertaTokenizer, RobertaModel
+from transformers import RobertaTokenizer, RobertaModel, BertTokenizer, BertModel
 from convlab.policy.vtrace_DPT.transformer_model.noisy_linear import NoisyLinear
 from convlab.policy.vtrace_DPT.create_descriptions import create_description_dicts
 
@@ -51,6 +51,11 @@ class NodeEmbedderRoberta(nn.Module):
                                                       f'embedded_descriptions_base_{self.dataset_name}.pt')
             if os.path.exists(embedded_descriptions_path):
                 self.embedded_descriptions = torch.load(embedded_descriptions_path).to(DEVICE)
+            else:
+                if dataset_name == "crosswoz":
+                    self.max_length = 40
+                    self.tokenizer = BertTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext")
+                    self.roberta_model = BertModel.from_pretrained("hfl/chinese-roberta-wwm-ext").to(DEVICE)
                 else:
                     self.tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
                     self.roberta_model = RobertaModel.from_pretrained("roberta-base").to(DEVICE)