Views
No views yet
1import torch
2import os
3from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
4from transformers import GenerationConfig, TextStreamer
5from peft import PeftModel
6from axolotl.prompters import AlpacaPrompter, PromptStyle
7
8### Load model
9torch_dtype = torch.bfloat16
10device_map = {"": int(os.environ.get("CUDA_DEVICE") or 0)}
11
12model_id = "nguyenthanhdo/noprob_model"
13
14tokenizer = LlamaTokenizer.from_pretrained(model_id)
15model = LlamaForCausalLM.from_pretrained(
16 model_id,
17 config=LlamaConfig.from_pretrained(model_id),
18 device_map=device_map,
19 torch_dtype=torch_dtype
20)
21
22### Build prompt
23prompter = AlpacaPrompter(prompt_style=PromptStyle.INSTRUCT.value)
24# instruction = "Provide short and concise answer. The answer should be straight and only provides explanation when needed." # Another instruction to test
25instruction = 'You are an AI assistant. Provide a detailed answer so user don’t need to search outside to understand the answer.'
26question = input()
27context = input()
28input = f"""Dựa vào bài viết dưới đây, trả lời câu hỏi phía dưới:\n{context}\n\nCâu hỏi: {question}"""
29prompt = prompter.build_prompt(instruction=instruction, input=input, output="").__next__()
30
31### Generate answer
32input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(model.device)
33model.eval()
34with torch.no_grad():
35 generation_config = GenerationConfig(
36 repetition_penalty=1.13,
37 max_new_tokens=max_new_tokens,
38 temperature=0.2,
39 top_p=0.95,
40 top_k=20,
41 pad_token_id=tokenizer.pad_token_id,
42 do_sample=True,
43 use_cache=True,
44 return_dict_in_generate=True,
45 output_attentions=False,
46 output_hidden_states=False,
47 output_scores=False,
48 )
49 streamer = TextStreamer(tokenizer, skip_prompt=True)
50 generated = model.generate(
51 inputs=input_ids,
52 generation_config=generation_config,
53 streamer=streamer,
54 )
55