Update modeling_quiet.py
Browse files- modeling_quiet.py +4 -4
modeling_quiet.py
CHANGED
|
@@ -1427,10 +1427,10 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
|
|
| 1427 |
@torch.no_grad()
|
| 1428 |
def generate(
|
| 1429 |
self,
|
| 1430 |
-
input_ids=None,
|
| 1431 |
-
attention_mask=None,
|
| 1432 |
-
max_new_tokens=None,
|
| 1433 |
-
temperature=1.1,
|
| 1434 |
**kwargs,
|
| 1435 |
):
|
| 1436 |
if isinstance(input_ids, str):
|
|
|
|
| 1427 |
@torch.no_grad()
|
| 1428 |
def generate(
|
| 1429 |
self,
|
| 1430 |
+
input_ids: Optional[torch.LongTensor] = None,
|
| 1431 |
+
attention_mask: Optional[torch.Tensor] = None,
|
| 1432 |
+
max_new_tokens: Optional[int] = None,
|
| 1433 |
+
temperature: float = 1.1,
|
| 1434 |
**kwargs,
|
| 1435 |
):
|
| 1436 |
if isinstance(input_ids, str):
|