aquiffoo commited on
Commit
c11d27f
·
verified ·
1 Parent(s): 7ce8e21

Update meshexpert.py

Browse files
Files changed (1) hide show
  1. meshexpert.py +16 -2
meshexpert.py CHANGED
@@ -1,3 +1,17 @@
1
- # Source code for MeshExpert from cell 39f2782d
2
- # Please replace this with the actual code from the notebook cell.
 
 
 
 
3
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig, PreTrainedModel, AutoModelForCausalLM # Import AutoModelForCausalLM
2
+ import torch
3
+ import torch.nn as nn
4
+ import torch.nn.functional as F
5
+ import math
6
+ from transformers.modeling_outputs import CausalLMOutputWithPast # Import the necessary output class
7
 
8
+ # Define a single Expert within the Mesh
9
+ class MeshExpert(nn.Module):
10
+ def __init__(self, config: MeshConfig):
11
+ super().__init__()
12
+ self.fc1 = nn.Linear(config.hidden_size, config.expert_intermediate_size)
13
+ self.gelu = nn.GELU() # Using GELU as an example activation
14
+ self.fc2 = nn.Linear(config.expert_intermediate_size, config.hidden_size)
15
+
16
+ def forward(self, x):
17
+ return self.fc2(self.gelu(self.fc1(x)))