Views
No views yet
1import torch
2from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
3from peft import PeftModel
4
5base_model="llma-7b"
6LORA_WEIGHTS = "llma-med-alpaca-7b"
7LOAD_8BIT = False
8
9tokenizer = LlamaTokenizer.from_pretrained(base_model)
10
11model = LlamaForCausalLM.from_pretrained(
12 base_model
13 load_in_8bit=LOAD_8BIT,
14 torch_dtype=torch.float16,
15 device_map="auto",
16)
17model = PeftModel.from_pretrained(
18 model,
19 LORA_WEIGHTS,
20 torch_dtype=torch.float16,
21)
22
23config = {
24 "temperature": 0 ,
25 "max_new_tokens": 1024,
26 "top_p": 0.5
27}
28
29prompt = "Translate to English: Je t’aime."
30input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
31outputs = model.generate(input_ids=input_ids, max_new_tokens=config["max_new_tokens"], temperature=config["temperature"])
32decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
33print(decoded[len(prompt):])
34
35