Update README.md
Browse files
README.md
CHANGED
@@ -22,11 +22,10 @@ The model, `Tevatron/dse-phi3-docmatix-v1.0`, is trained using the `Tevatron/doc
|
|
22 |
|
23 |
```python
|
24 |
import torch
|
25 |
-
from transformers import AutoProcessor, AutoModelForCausalLM
|
26 |
|
27 |
-
processor = AutoProcessor.from_pretrained('
|
28 |
-
|
29 |
-
model = AutoModelForCausalLM.from_pretrained('Tevatron/dse-phi3-docmatix-v1.0', trust_remote_code=True, config=config, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16).to('cuda:0')
|
30 |
|
31 |
def get_embedding(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
|
32 |
sequence_lengths = attention_mask.sum(dim=1) - 1
|
|
|
22 |
|
23 |
```python
|
24 |
import torch
|
25 |
+
from transformers import AutoProcessor, AutoModelForCausalLM
|
26 |
|
27 |
+
processor = AutoProcessor.from_pretrained('Tevatron/dse-phi3-docmatix-v1.0', trust_remote_code=True)
|
28 |
+
model = AutoModelForCausalLM.from_pretrained('Tevatron/dse-phi3-docmatix-v1.0', trust_remote_code=True, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, use_cache=False).to('cuda:0')
|
|
|
29 |
|
30 |
def get_embedding(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
|
31 |
sequence_lengths = attention_mask.sum(dim=1) - 1
|