bigdefence
commited on
Commit
•
177bfea
1
Parent(s):
0ccc156
Update README.md
Browse files
README.md
CHANGED
@@ -11,6 +11,42 @@ tags:
|
|
11 |
- trl
|
12 |
---
|
13 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
14 |
# Uploaded model
|
15 |
|
16 |
- **Developed by:** bigdefence
|
|
|
11 |
- trl
|
12 |
---
|
13 |
|
14 |
+
```
|
15 |
+
import torch
|
16 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
17 |
+
base_model = 'bigdefence/Llama-3.1-8B-bigdefence-Ko'
|
18 |
+
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
19 |
+
|
20 |
+
tokenizer = AutoTokenizer.from_pretrained(base_model)
|
21 |
+
model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.float16, device_map="auto")
|
22 |
+
model.eval()
|
23 |
+
def generate_response(prompt, model, tokenizer, text_streamer,max_new_tokens=256):
|
24 |
+
inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
|
25 |
+
inputs = inputs.to(model.device)
|
26 |
+
|
27 |
+
with torch.no_grad():
|
28 |
+
outputs = model.generate(
|
29 |
+
**inputs,
|
30 |
+
streamer=text_streamer,
|
31 |
+
max_new_tokens=max_new_tokens,
|
32 |
+
do_sample=True,
|
33 |
+
pad_token_id=tokenizer.eos_token_id
|
34 |
+
)
|
35 |
+
|
36 |
+
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
37 |
+
return response.replace(prompt, '').strip()
|
38 |
+
key = "안녕?"
|
39 |
+
prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
|
40 |
+
|
41 |
+
### Instruction:
|
42 |
+
{key}
|
43 |
+
|
44 |
+
### Response:
|
45 |
+
"""
|
46 |
+
text_streamer = TextStreamer(tokenizer)
|
47 |
+
response = generate_response(prompt, model, tokenizer,text_streamer)
|
48 |
+
print(response)
|
49 |
+
```
|
50 |
# Uploaded model
|
51 |
|
52 |
- **Developed by:** bigdefence
|