huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-10 11:54:52 +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 import torch
from abc import ABC, abstractmethod 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 transformers import PreTrainedTokenizerBase
from text_generation_server.models.types import Batch, GeneratedText 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) B = TypeVar("B", bound=Batch)
class Model(ABC): class Model(ABC, Generic[B]):
def __init__( def __init__(
self, self,
model: torch.nn.Module, model: torch.nn.Module,