File size: 2,451 Bytes
430a3fd
7cb041b
e5c2e8f
 
 
 
 
 
 
 
 
 
 
477a05d
 
 
 
87f98f5
 
 
 
 
477a05d
 
 
 
 
 
 
edc8f54
eff2613
 
 
 
 
 
 
 
 
 
 
87fd4ce
a7b00bc
eff2613
 
a7b00bc
 
 
 
87f98f5
 
 
a7b00bc
eff2613
477a05d
a7b00bc
 
 
 
edc8f54
 
a7b00bc
 
edc8f54
 
 
477a05d
67629f0
 
477a05d
acad2d2
 
 
99f6c76
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
---
base_model: AI-Sweden-Models/gpt-sw3-6.7b-v2-instruct
language:
- sv
- da
- 'no'
- en
pipeline_tag: text-generation
inference:
  parameters:
    temperature: 0.7
tags:
- translation
---
# Model Card for gpt-sw3-6.7b-v2-translator
The `gpt-sw3-6.7b-v2-translator` is a finetuned version of `gpt-sw3-6.7b-v2-instruct` on a carefully selected translation pair dataset that was gathered by AI Sweden.


## Intended usage:
Translate text data from English to Swedish, or Swedish to English.


## How to use:
```python
import torch
from transformers import pipeline, StoppingCriteriaList, StoppingCriteria

device = "cuda" if torch.cuda.is_available() else "cpu"


# (Optional) - define a stopping criteria
# We ideally want the model to stop generate once the response from the Bot is generated
class StopOnTokenCriteria(StoppingCriteria):
    def __init__(self, stop_token_id):
        self.stop_token_id = stop_token_id

    def __call__(self, input_ids, scores, **kwargs):
        return input_ids[0, -1] == self.stop_token_id


pipe = pipeline(
    task="text-generation",
    model="AI-Sweden-Models/gpt-sw3-6.7b-v2-translator",
    device=device
)

stop_on_token_criteria = StopOnTokenCriteria(stop_token_id=pipe.tokenizer.bos_token_id)
text = "I like to eat ice cream in the summer."

# This will translate English to Swedish
# To translate from Swedish to English the prompt would be:
# prompt = f"<|endoftext|><s>User: Översätt till Engelska från Svenska\n{text}<s>Bot:"

prompt = f"<|endoftext|><s>User: Översätt till Svenska från Engelska\n{text}<s>Bot:"

input_tokens = pipe.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
max_model_length = 2048
dynamic_max_length = max_model_length - input_tokens.shape[1]

response = pipe(
    prompt,
    max_length=dynamic_max_length,
    truncation=True,
    stopping_criteria=StoppingCriteriaList([stop_on_token_criteria])
)

print(response[0]["generated_text"].split("<s>Bot: ")[-1])
```
```python
>>> "Jag tycker om att äta glass på sommaren."
```

## Training & Data:
The training was done on 1 NVIDIA DGX using DeepSpeed ZeRO 3 for three epochs on roughly 4GB of carefully selected translation data. It is a full finetune of all of the model parameters. 

| Epoch | Training Loss | Evaluation Loss |
|-------|---------------|-----------------|
| 1     | 1.309         | 1.281           |
| 2     | 1.161         | 1.242           |
| 3     | 1.053         | 1.219           |