Views
No views yet

1import torch
2from transformers import pipeline
3from transformers import TextStreamer
4from transformers import AutoTokenizer
5from transformers import AutoModelForCausalLM
6
7
8model_name = 'four-two-labs/lynx-micro'
9
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11model = AutoModelForCausalLM.from_pretrained(
12 model_name,
13 device_map='cuda',
14 torch_dtype=torch.bfloat16,
15 use_flash_attention_2=True, # Remove if flash attention isn't available
16)
17
18pipe = pipeline(
19 'text-generation',
20 model=model,
21 tokenizer=tokenizer,
22 streamer=TextStreamer(tokenizer=tokenizer)
23)
24
25messages = [
26 #{'role': 'user', 'content': 'Lös ekvationen 2x^2-5 = 9'},
27 #{'role': 'user', 'content': 'Vad är fel med denna mening: "Hej! Jag idag bra mår."'},
28 #{'role': 'user', 'content': """Översätt till svenska: Hayashi, the Japanese government spokesperson, said Monday that Tokyo is answering the Chinese presence around the islands with vessels of its own.\n\n“We ensure a comprehensive security system for territorial waters by deploying Coast Guard patrol vessels that are consistently superior to other party’s capacity,” Hayashi said.\n\nAny Japanese-Chinese incident in the Senkakus raises the risk of a wider conflict, analysts note, due to Japan’s mutual defense treaty with the United States.\n\nWashington has made clear on numerous occasions that it considers the Senkakus to be covered by the mutual defense pact."""},
29 #{'role': 'user', 'content': """Vad handlar texten om?\n\nHayashi, the Japanese government spokesperson, said Monday that Tokyo is answering the Chinese presence around the islands with vessels of its own.\n\n“We ensure a comprehensive security system for territorial waters by deploying Coast Guard patrol vessels that are consistently superior to other party’s capacity,” Hayashi said.\n\nAny Japanese-Chinese incident in the Senkakus raises the risk of a wider conflict, analysts note, due to Japan’s mutual defense treaty with the United States.\n\nWashington has made clear on numerous occasions that it considers the Senkakus to be covered by the mutual defense pact."""},
30 #{'role': 'user', 'content': """Skriv en sci-fi novell som utspelar sig över millenium på en planet runt ett binärt stjärnsystem."""},
31 {'role': 'user', 'content': 'Hur många helikoptrar kan en människa äta på en gång?'},
32]
33
34r = pipe(
35 messages,
36 max_length=4096,
37 do_sample=False,
38 eos_token_id=[tokenizer.vocab['<end_of_turn>'], tokenizer.eos_token_id],
39)
