Unverified Commit 85295621 authored by Suraj Patil's avatar Suraj Patil Committed by GitHub
Browse files

Fix blenderbot conversion script (#16472)

parent c85547af
...@@ -18,7 +18,7 @@ import argparse ...@@ -18,7 +18,7 @@ import argparse
import torch import torch
from transformers import BartConfig, BartForConditionalGeneration from transformers import BlenderbotConfig, BlenderbotForConditionalGeneration
from transformers.utils import logging from transformers.utils import logging
...@@ -81,8 +81,8 @@ def convert_parlai_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_ ...@@ -81,8 +81,8 @@ def convert_parlai_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_
""" """
model = torch.load(checkpoint_path, map_location="cpu") model = torch.load(checkpoint_path, map_location="cpu")
sd = model["model"] sd = model["model"]
cfg = BartConfig.from_json_file(config_json_path) cfg = BlenderbotConfig.from_json_file(config_json_path)
m = BartForConditionalGeneration(cfg) m = BlenderbotForConditionalGeneration(cfg)
valid_keys = m.model.state_dict().keys() valid_keys = m.model.state_dict().keys()
failures = [] failures = []
mapping = {} mapping = {}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment