Update services/model_service.py
Browse files
services/model_service.py
CHANGED
@@ -44,6 +44,7 @@ class ModelService:
|
|
44 |
# Load model with the updated configuration
|
45 |
self.model = AutoModelForCausalLM.from_pretrained(
|
46 |
settings.MODEL_NAME,
|
|
|
47 |
torch_dtype=torch.float16 if settings.DEVICE == "cuda" else torch.float32,
|
48 |
device_map="auto" if settings.DEVICE == "cuda" else None,
|
49 |
config=config
|
|
|
44 |
# Load model with the updated configuration
|
45 |
self.model = AutoModelForCausalLM.from_pretrained(
|
46 |
settings.MODEL_NAME,
|
47 |
+
model_type = "llama",
|
48 |
torch_dtype=torch.float16 if settings.DEVICE == "cuda" else torch.float32,
|
49 |
device_map="auto" if settings.DEVICE == "cuda" else None,
|
50 |
config=config
|