Coool2 commited on
Commit
5fa5125
·
1 Parent(s): 3c84018

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +4 -4
agent.py CHANGED
@@ -44,7 +44,7 @@ from llama_index.readers.file import (
44
  PandasExcelReader,
45
  VideoAudioReader # Adding VideoAudioReader for handling audio/video without API
46
  )
47
- from pydantic import Field
48
 
49
  # Optional API-based imports (conditionally loaded)
50
  try:
@@ -131,9 +131,9 @@ def initialize_models(use_api_mode=False):
131
  context_window: int = 32768
132
  num_output: int = 256
133
  model_name: str = "Qwen/Qwen2.5-VL-7B-Instruct"
134
-
135
- def __init__(self):
136
- self.device = "cuda"
137
  self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
138
  self.model_name, torch_dtype=torch.bfloat16, device_map="auto"
139
  )
 
44
  PandasExcelReader,
45
  VideoAudioReader # Adding VideoAudioReader for handling audio/video without API
46
  )
47
+ from pydantic import PrivateAttr
48
 
49
  # Optional API-based imports (conditionally loaded)
50
  try:
 
131
  context_window: int = 32768
132
  num_output: int = 256
133
  model_name: str = "Qwen/Qwen2.5-VL-7B-Instruct"
134
+ _device: str = PrivateAttr()
135
+ def __init__(self, device: str = "cuda", **kwargs):
136
+ self._device = "cuda"
137
  self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
138
  self.model_name, torch_dtype=torch.bfloat16, device_map="auto"
139
  )