chore: Remove commented out code in load_gpu_model function
Browse files- kitt/core/model.py +2 -2
kitt/core/model.py
CHANGED
|
@@ -6,13 +6,13 @@ import xml.etree.ElementTree as ET
|
|
| 6 |
from enum import Enum
|
| 7 |
from typing import List
|
| 8 |
|
|
|
|
| 9 |
from langchain.tools.base import StructuredTool
|
| 10 |
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
| 11 |
from langchain_core.utils.function_calling import convert_to_openai_tool
|
| 12 |
from loguru import logger
|
| 13 |
from ollama import Client
|
| 14 |
from pydantic import BaseModel
|
| 15 |
-
import torch
|
| 16 |
|
| 17 |
from kitt.skills import vehicle_status
|
| 18 |
from kitt.skills.common import config
|
|
@@ -358,7 +358,7 @@ def load_gpu_model():
|
|
| 358 |
device_map="auto",
|
| 359 |
load_in_8bit=False,
|
| 360 |
load_in_4bit=True,
|
| 361 |
-
use_flash_attention_2=True,
|
| 362 |
)
|
| 363 |
|
| 364 |
return model, tokenizer
|
|
|
|
| 6 |
from enum import Enum
|
| 7 |
from typing import List
|
| 8 |
|
| 9 |
+
import torch
|
| 10 |
from langchain.tools.base import StructuredTool
|
| 11 |
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
| 12 |
from langchain_core.utils.function_calling import convert_to_openai_tool
|
| 13 |
from loguru import logger
|
| 14 |
from ollama import Client
|
| 15 |
from pydantic import BaseModel
|
|
|
|
| 16 |
|
| 17 |
from kitt.skills import vehicle_status
|
| 18 |
from kitt.skills.common import config
|
|
|
|
| 358 |
device_map="auto",
|
| 359 |
load_in_8bit=False,
|
| 360 |
load_in_4bit=True,
|
| 361 |
+
# use_flash_attention_2=True,
|
| 362 |
)
|
| 363 |
|
| 364 |
return model, tokenizer
|