Image-Text-to-Text
sentence-transformers
Safetensors
Transformers
qwen2_vl
Qwen2-VL
conversational
cheesyFishes commited on
Commit
c43084e
·
verified ·
1 Parent(s): 0da786d

cpu default, fix args

Browse files
Files changed (1) hide show
  1. custom_st.py +8 -19
custom_st.py CHANGED
@@ -24,8 +24,10 @@ class Transformer(nn.Module):
24
  max_seq_length: Optional[int] = None,
25
  model_args: Optional[Dict[str, Any]] = None,
26
  processor_args: Optional[Dict[str, Any]] = None,
 
 
27
  cache_dir: Optional[str] = None,
28
- device: str = 'cuda:0',
29
  backend: Literal['torch', 'onnx', 'openvino'] = 'torch',
30
  **kwargs,
31
  ) -> None:
@@ -54,24 +56,11 @@ class Transformer(nn.Module):
54
  })
55
 
56
  # Initialize model
57
- try:
58
- self.model = Qwen2VLForConditionalGeneration.from_pretrained(
59
- model_name_or_path,
60
- attn_implementation="flash_attention_2",
61
- torch_dtype=torch.bfloat16,
62
- device_map=device,
63
- cache_dir=cache_dir,
64
- **model_kwargs
65
- ).eval()
66
- except (ImportError, ValueError) as e:
67
- print(f"Flash attention not available, falling back to default attention: {e}")
68
- self.model = Qwen2VLForConditionalGeneration.from_pretrained(
69
- model_name_or_path,
70
- torch_dtype=torch.bfloat16,
71
- device_map=device,
72
- cache_dir=cache_dir,
73
- **model_kwargs
74
- ).eval()
75
 
76
  # Initialize processor
77
  self.processor = AutoProcessor.from_pretrained(
 
24
  max_seq_length: Optional[int] = None,
25
  model_args: Optional[Dict[str, Any]] = None,
26
  processor_args: Optional[Dict[str, Any]] = None,
27
+ tokenizer_args: Optional[Dict[str, Any]] = None,
28
+ config_args: Optional[Dict[str, Any]] = None,
29
  cache_dir: Optional[str] = None,
30
+ device: str = 'cpu',
31
  backend: Literal['torch', 'onnx', 'openvino'] = 'torch',
32
  **kwargs,
33
  ) -> None:
 
56
  })
57
 
58
  # Initialize model
59
+ self.model = Qwen2VLForConditionalGeneration.from_pretrained(
60
+ model_name_or_path,
61
+ cache_dir=cache_dir,
62
+ **model_kwargs
63
+ ).eval()
 
 
 
 
 
 
 
 
 
 
 
 
 
64
 
65
  # Initialize processor
66
  self.processor = AutoProcessor.from_pretrained(