change nn to models (#3032)

applications/ChatGPT/chatgpt/models/opt/__init__.py

+from .opt_actor import OPTActor
+from .opt_critic import OPTCritic
+from .opt_rm import OPTRM
+
+__all__ = ['OPTActor', 'OPTCritic', 'OPTRM']

applications/ChatGPT/chatgpt/nn/opt_actor.py → applications/ChatGPT/chatgpt/models/opt/opt_actor.py

@@ -3,7 +3,7 @@ from typing import Optional
 from transformers.models.opt.configuration_opt import OPTConfig
 from transformers.models.opt.modeling_opt import OPTForCausalLM
 
-from .actor import Actor
+from ..base import Actor
 
 
 class OPTActor(Actor):

applications/ChatGPT/chatgpt/nn/opt_critic.py → applications/ChatGPT/chatgpt/models/opt/opt_critic.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from transformers.models.opt.configuration_opt import OPTConfig
 from transformers.models.opt.modeling_opt import OPTModel
 
-from .critic import Critic
+from ..base import Critic
 
 
 class OPTCritic(Critic):

applications/ChatGPT/chatgpt/nn/opt_rm.py → applications/ChatGPT/chatgpt/models/opt/opt_rm.py

@@ -3,7 +3,7 @@ from typing import Optional
 import torch.nn as nn
 from transformers import OPTConfig, OPTModel
 
-from .reward_model import RewardModel
+from ..base import RewardModel
 
 
 class OPTRM(RewardModel):

applications/ChatGPT/chatgpt/nn/__init__.py

-from .actor import Actor
-from .bloom_actor import BLOOMActor
-from .bloom_critic import BLOOMCritic
-from .bloom_rm import BLOOMRM
-from .critic import Critic
-from .gpt_actor import GPTActor
-from .gpt_critic import GPTCritic
-from .gpt_rm import GPTRM
-from .loss import PairWiseLoss, PolicyLoss, PPOPtxActorLoss, ValueLoss
-from .opt_actor import OPTActor
-from .opt_critic import OPTCritic
-from .opt_rm import OPTRM
-from .reward_model import RewardModel
-
-__all__ = [
-    'Actor', 'Critic', 'RewardModel', 'PolicyLoss', 'ValueLoss', 'PPOPtxActorLoss', 'PairWiseLoss', 'GPTActor',
-    'GPTCritic', 'GPTRM', 'BLOOMActor', 'BLOOMCritic', 'BLOOMRM', 'OPTActor', 'OPTCritic', 'OPTRM'
-]

applications/ChatGPT/chatgpt/trainer/ppo.py

@@ -2,8 +2,9 @@ from typing import Any, Callable, Dict, List, Optional
 
 import torch.nn as nn
 from chatgpt.experience_maker import Experience, NaiveExperienceMaker
-from chatgpt.nn import Actor, Critic, PolicyLoss, ValueLoss
-from chatgpt.nn.generation_utils import update_model_kwargs_fn
+from chatgpt.models.base import Actor, Critic
+from chatgpt.models.generation_utils import update_model_kwargs_fn
+from chatgpt.models.loss import PolicyLoss, ValueLoss
 from chatgpt.replay_buffer import NaiveReplayBuffer
 from torch.optim import Optimizer
 

applications/ChatGPT/chatgpt/trainer/rm.py

@@ -3,7 +3,7 @@ from abc import ABC
 import loralib as lora
 import torch
 from chatgpt.dataset import RewardDataset
-from chatgpt.nn import PairWiseLoss
+from chatgpt.models.loss import PairWiseLoss
 from torch.optim import Adam, Optimizer
 from torch.utils.data import DataLoader
 from tqdm import tqdm

applications/ChatGPT/chatgpt/trainer/strategies/base.py

@@ -5,7 +5,7 @@ from typing import Any, List, Tuple, Union
 import numpy as np
 import torch
 import torch.nn as nn
-from chatgpt.nn import Actor
+from chatgpt.models.base import Actor, Critic, RewardModel
 from chatgpt.replay_buffer import ReplayBuffer
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader

applications/ChatGPT/chatgpt/trainer/strategies/colossalai.py

@@ -5,7 +5,7 @@ import torch
 import torch.distributed as dist
 import torch.nn as nn
 import torch.optim as optim
-from chatgpt.nn import Actor
+from chatgpt.models.base import Actor
 from torch.optim import Optimizer
 
 import colossalai

applications/ChatGPT/chatgpt/trainer/strategies/ddp.py

@@ -5,7 +5,7 @@ import numpy as np
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from chatgpt.nn import Actor
+from chatgpt.models.base import Actor
 from chatgpt.replay_buffer import ReplayBuffer
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer

applications/ChatGPT/examples/inference.py

 import argparse
 
 import torch
-from chatgpt.nn import BLOOMActor, GPTActor, OPTActor
+from chatgpt.models.bloom import BLOOMActor
+from chatgpt.models.gpt import GPTActor
+from chatgpt.models.opt import OPTActor
 from transformers import AutoTokenizer
 from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
 

applications/ChatGPT/examples/train_dummy.py

@@ -2,7 +2,10 @@ import argparse
 from copy import deepcopy
 
 import torch
-from chatgpt.nn import BLOOMActor, BLOOMCritic, GPTActor, GPTCritic, OPTActor, OPTCritic, RewardModel
+from chatgpt.models.base import RewardModel
+from chatgpt.models.bloom import BLOOMActor, BLOOMCritic
+from chatgpt.models.gpt import GPTActor, GPTCritic
+from chatgpt.models.opt import OPTActor, OPTCritic
 from chatgpt.trainer import PPOTrainer
 from chatgpt.trainer.callbacks import SaveCheckpoint
 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy

applications/ChatGPT/examples/train_prompts.py

@@ -3,7 +3,10 @@ from copy import deepcopy
 
 import pandas as pd
 import torch
-from chatgpt.nn import BLOOMActor, BLOOMCritic, GPTActor, GPTCritic, OPTActor, OPTCritic, RewardModel
+from chatgpt.models.base import RewardModel
+from chatgpt.models.bloom import BLOOMActor, BLOOMCritic
+from chatgpt.models.gpt import GPTActor, GPTCritic
+from chatgpt.models.opt import OPTActor, OPTCritic
 from chatgpt.trainer import PPOTrainer
 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
 from torch.optim import Adam

applications/ChatGPT/examples/train_reward_model.py

@@ -3,7 +3,10 @@ import argparse
 import loralib as lora
 import torch
 from chatgpt.dataset import RewardDataset
-from chatgpt.nn import BLOOMRM, GPTRM, OPTRM
+from chatgpt.models.base import RewardModel
+from chatgpt.models.bloom import BLOOMRM
+from chatgpt.models.gpt import GPTRM
+from chatgpt.models.opt import OPTRM
 from chatgpt.trainer import RewardModelTrainer
 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy
 from datasets import load_dataset

applications/ChatGPT/tests/test_checkpoint.py

@@ -7,7 +7,7 @@ import pytest
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
-from chatgpt.nn import GPTActor
+from chatgpt.models.gpt import GPTActor
 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 

applications/ChatGPT/tests/test_data.py

@@ -7,7 +7,8 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 from chatgpt.experience_maker import NaiveExperienceMaker
-from chatgpt.nn import GPTActor, GPTCritic, RewardModel
+from chatgpt.models.base import RewardModel
+from chatgpt.models.gpt import GPTActor, GPTCritic
 from chatgpt.replay_buffer import NaiveReplayBuffer
 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config