fix params

This commit is contained in:
josc146
2023-07-28 22:13:19 +08:00
parent d2928e54f7
commit 9d89b6f4db
3 changed files with 7 additions and 7 deletions

View File

@@ -3,7 +3,7 @@ import os
import pathlib
import copy
import re
from typing import Dict, Iterable, List, Tuple
from typing import Dict, Iterable, List, Tuple, Union
from utils.log import quick_log
from fastapi import HTTPException
from pydantic import BaseModel, Field
@@ -213,7 +213,7 @@ class AbstractRWKV(ABC):
return state[0].tolist(), token_len
def generate(
self, prompt: str, stop: str | List[str] = None
self, prompt: str, stop: Union[str, List[str]] = None
) -> Iterable[Tuple[str, str, int, int]]:
quick_log(None, None, "Generation Prompt:\n" + prompt)
cache = None