@@ -41,17 +41,28 @@ class RetrieverConfig(ChangesMixin):
4141 The name of the retriever model to use. It must be a HuggingFace repo id.
4242 retriever_type: str
4343 The type of retriever to use.
44+ knowledge_base: KnowledgeBase
45+ The knowledge base to use for the retriever.
46+ index_status: str
47+ The status of the retriever index.
48+ s3_index_path: str
49+ The path to the retriever index in S3.
4450 batch_size: int
4551 The batch size to use for the retriever.
4652 device: str
4753 The device to use for the retriever.
54+ enabled: bool
55+ Whether the retriever is enabled.
56+ num_replicas: int
57+ The number of replicas to deploy in the Ray cluster.
4858 """
4959
5060 objects = models .Manager () # The default manager.
5161
5262 enabled_objects = EnabledRetrieverConfigManager () # The Dahl-specific manager.
53-
5463 name = models .CharField (max_length = 255 , unique = True )
64+
65+ # Model properties
5566 model_name = models .CharField (
5667 max_length = 255 , default = "colbert-ir/colbertv2.0"
5768 ) # For dev and demo purposes.
@@ -60,13 +71,8 @@ class RetrieverConfig(ChangesMixin):
6071 choices = RetrieverTypeChoices .choices ,
6172 default = RetrieverTypeChoices .COLBERT ,
6273 )
63- batch_size = models .IntegerField (
64- default = 1
65- ) # batch size 1 for better default cpu generation
66- device = models .CharField (
67- max_length = 10 , choices = DeviceChoices .choices , default = DeviceChoices .CPU
68- )
6974
75+ # Knowledge Base properties
7076 knowledge_base = models .ForeignKey (KnowledgeBase , on_delete = models .CASCADE )
7177 index_status = models .CharField (
7278 max_length = 20 ,
@@ -78,6 +84,13 @@ class RetrieverConfig(ChangesMixin):
7884 max_length = 255 , blank = True , null = True , editable = False
7985 )
8086
87+ # Model inference properties
88+ batch_size = models .IntegerField (
89+ default = 1
90+ ) # batch size 1 for better default cpu generation
91+ device = models .CharField (
92+ max_length = 10 , choices = DeviceChoices .choices , default = DeviceChoices .CPU
93+ )
8194 enabled = models .BooleanField (default = False )
8295 num_replicas = models .IntegerField (default = 1 )
8396
@@ -233,6 +246,10 @@ class LLMConfig(ChangesMixin):
233246 The base url where the model is hosted. It is used for vLLM deployments and Together LLM Endpoints.
234247 model_max_length: int
235248 The maximum length of the model.
249+ enabled: bool
250+ Whether the LLM is enabled.
251+ num_replicas: int
252+ The number of replicas to deploy in the Ray cluster.
236253 """
237254
238255 objects = models .Manager () # The default manager.
0 commit comments