Upload ModularStarEncoder
Browse files- modularStarEncoder.py +2 -3
modularStarEncoder.py
CHANGED
@@ -1,6 +1,6 @@
|
|
1 |
from transformers import Starcoder2Model
|
2 |
import sys
|
3 |
-
from
|
4 |
import os
|
5 |
from dataclasses import dataclass
|
6 |
from typing import Optional, Tuple, Union, List
|
@@ -298,7 +298,6 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
|
|
298 |
|
299 |
|
300 |
"""
|
301 |
-
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
302 |
|
303 |
source_embedding = self.starEncoder2(
|
304 |
input_ids,
|
@@ -307,7 +306,7 @@ class ModularStarEncoder(StarEncoder2PreTrainedModel):
|
|
307 |
inputs_embeds=inputs_embeds,
|
308 |
output_attentions=output_attentions,
|
309 |
output_hidden_states=True,
|
310 |
-
return_dict=
|
311 |
).hidden_states
|
312 |
|
313 |
|
|
|
1 |
from transformers import Starcoder2Model
|
2 |
import sys
|
3 |
+
from config import ModularStarEncoderConfig
|
4 |
import os
|
5 |
from dataclasses import dataclass
|
6 |
from typing import Optional, Tuple, Union, List
|
|
|
298 |
|
299 |
|
300 |
"""
|
|
|
301 |
|
302 |
source_embedding = self.starEncoder2(
|
303 |
input_ids,
|
|
|
306 |
inputs_embeds=inputs_embeds,
|
307 |
output_attentions=output_attentions,
|
308 |
output_hidden_states=True,
|
309 |
+
return_dict=True,
|
310 |
).hidden_states
|
311 |
|
312 |
|