Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3tokenizer = AutoTokenizer.from_pretrained("vlsp-2023-vllm/hoa-1b4")
4model = AutoModelForCausalLM.from_pretrained("vlsp-2023-vllm/hoa-1b4", low_cpu_mem_usage=True)
5
6device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
7model.to(device)
8
9prompt = "Địa chỉ trường Đại học Tôn Đức Thắng nằm ở số"
10input_ids = tokenizer(prompt, return_tensors="pt")['input_ids'].to(device)
11
12gen_tokens = model.generate(input_ids, max_length=max_length, repetition_penalty=1.1)
13
14print(tokenizer.batch_decode(gen_tokens)[0])