huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-09 19:34:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make B a generic type of Model

Browse Source 
This PR fixes typing error at `def batch_type`
This commit is contained in:
Yang, Bo 2023-05-18 07:29:20 -07:00 committed by GitHub
parent 5a58226130
commit 27d30f685a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
server/text_generation_server/models/model.py
View File

@ -1,7 +1,7 @@
import torch
from abc import ABC, abstractmethod
from typing import List, Tuple, Optional, TypeVar, Type
from typing import Generic, List, Tuple, Optional, TypeVar, Type
from transformers import PreTrainedTokenizerBase
from text_generation_server.models.types import Batch, GeneratedText
@ -10,7 +10,7 @@ from text_generation_server.pb.generate_pb2 import InfoResponse
B = TypeVar("B", bound=Batch)
class Model(ABC):
class Model(ABC, Generic[B]):
def __init__(
self,
model: torch.nn.Module,