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

fix: improve text model loading

Browse Source
This commit is contained in:
drbh 2024-12-18 14:58:27 +00:00
parent 34174af8c8
commit 064e040ee3
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server/text_generation_server/models/custom_modeling/idefics2.py
View File

@ -723,7 +723,7 @@ class Idefics3ForConditionalGeneration(nn.Module):
vision_config = config.vision_config vision_config = config.vision_config
self.text_model = load_text_model( self.text_model = load_text_model(
prefix="model" if not prefix else f"{prefix}.model", prefix=f"{prefix}.model.text_model" if prefix else "model.text_model",
config=config.text_config, config=config.text_config,
weights=weights, weights=weights,
name="text_model", name="text_model",

2
server/text_generation_server/models/custom_modeling/vlm.py
View File

@ -4,7 +4,7 @@ def load_text_model(prefix, config, weights, name=None):
FlashLlamaForCausalLM, FlashLlamaForCausalLM,
) )
return FlashLlamaForCausalLM(f"{prefix}.text_model", config, weights) return FlashLlamaForCausalLM(prefix, config, weights)
elif config.model_type == "mistral": elif config.model_type == "mistral":
from text_generation_server.models.custom_modeling.flash_mistral_modeling import ( from text_generation_server.models.custom_modeling.flash_mistral_modeling import (
FlashMistralForCausalLM, FlashMistralForCausalLM,