ameliechatelain commited on
Commit
c25e8b6
·
verified ·
1 Parent(s): 48767d3

Model card: add exact reranking prompts, mark flash-attn optional

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -116,7 +116,7 @@ model_id = "lightonai/LightOn-rerank-LW-4B"
116
  model = AutoModelForImageTextToText.from_pretrained(
117
  model_id,
118
  dtype=torch.bfloat16,
119
- attn_implementation="flash_attention_2",
120
  device_map="cuda",
121
  ).eval()
122
  processor = AutoProcessor.from_pretrained(model_id)
 
116
  model = AutoModelForImageTextToText.from_pretrained(
117
  model_id,
118
  dtype=torch.bfloat16,
119
+ attn_implementation="flash_attention_2", # optional — remove if flash-attn is not installed
120
  device_map="cuda",
121
  ).eval()
122
  processor = AutoProcessor.from_pretrained(model_id)