Update ChatFalcon.py
Browse files- ChatFalcon.py +5 -1
ChatFalcon.py
CHANGED
@@ -22,7 +22,7 @@ import torch.utils.checkpoint
|
|
22 |
from torch import nn
|
23 |
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
|
24 |
from torch.nn import functional as F
|
25 |
-
|
26 |
from transformers.modeling_outputs import (
|
27 |
BaseModelOutputWithPastAndCrossAttentions,
|
28 |
CausalLMOutputWithCrossAttentions,
|
@@ -1265,3 +1265,7 @@ class FalconForQuestionAnswering(FalconPreTrainedModel):
|
|
1265 |
hidden_states=outputs.hidden_states,
|
1266 |
attentions=outputs.attentions,
|
1267 |
)
|
|
|
|
|
|
|
|
|
|
22 |
from torch import nn
|
23 |
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
|
24 |
from torch.nn import functional as F
|
25 |
+
from transformers import AutoModel,AutoModelForCausalLM, AutoConfig
|
26 |
from transformers.modeling_outputs import (
|
27 |
BaseModelOutputWithPastAndCrossAttentions,
|
28 |
CausalLMOutputWithCrossAttentions,
|
|
|
1265 |
hidden_states=outputs.hidden_states,
|
1266 |
attentions=outputs.attentions,
|
1267 |
)
|
1268 |
+
|
1269 |
+
AutoConfig.register("falcon", FalconConfig)
|
1270 |
+
AutoModel.register(FalconConfig, FalconModel)
|
1271 |
+
AutoModelForCausalLM.register(FalconConfig, FalconForCausalLM)
|