Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
2 |
+
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
3 |
+
import torch
|
4 |
+
|
5 |
+
|
6 |
+
# Cấu hình mô hình
|
7 |
+
MODEL = "Viet-Mistral/Vistral-7B-Chat"
|
8 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
9 |
+
print('device =', device)
|
10 |
+
|
11 |
+
# Load mô hình và tokenizer
|
12 |
+
model = AutoModelForCausalLM.from_pretrained(
|
13 |
+
'Viet-Mistral/Vistral-7B-Chat',
|
14 |
+
torch_dtype=torch.bfloat16, # change to torch.float16 if you're using V100
|
15 |
+
device_map="auto",
|
16 |
+
use_cache=True,
|
17 |
+
#cache_dir='/workspace/thviet/hf_cache'
|
18 |
+
)
|
19 |
+
|
20 |
+
tokenizer = AutoTokenizer.from_pretrained(MODEL, cache_dir='/workspace/thviet/hf_cache')
|
21 |
+
|
22 |
+
lora_config = LoraConfig.from_pretrained(
|
23 |
+
"thviet79/model-QA-medical", # Thay bằng đường dẫn đến mô hình LoRA trên Hugging Face
|
24 |
+
#cache_dir='/workspace/thviet/hf_cache'
|
25 |
+
)
|
26 |
+
|
27 |
+
# Áp dụng cấu hình LoRA vào mô hình
|
28 |
+
model = get_peft_model(model, lora_config)
|
29 |
+
# Chuẩn bị hội thoại và input
|
30 |
+
system_prompt = "Bạn là một trợ lí ảo Tiếng Việt về lĩnh vực y tế."
|
31 |
+
question = "Chào bác sĩ,\nRăng cháu hiện tại có mủ ở dưới lợi nhưng khi đau cháu sẽ không ngủ được (quá đau). Tuy nhiên chỉ vài ngày là hết mà thỉnh thoảng nó lại bị đau. Chị cháu bảo là trước chị cháu cũng bị như vậy chỉ là đau răng tuổi dậy thì thôi. Bác sĩ cho cháu hỏi đau răng kèm có mủ dưới lợi là bệnh gì? Cháu có cần đi chữa trị không? Cháu cảm ơn."
|
32 |
+
|
33 |
+
conversation = [{"role": "system", "content": system_prompt }]
|
34 |
+
human = f"Vui lòng trả lời câu hỏi sau: {question}"
|
35 |
+
conversation.append({"role": "user", "content": human })
|
36 |
+
|
37 |
+
# Chuyển các tensor đầu vào sang đúng thiết bị
|
38 |
+
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(device)
|
39 |
+
|
40 |
+
# Tạo đầu ra từ mô hình
|
41 |
+
out_ids = model.generate(
|
42 |
+
input_ids=input_ids,
|
43 |
+
max_new_tokens=768,
|
44 |
+
do_sample=True,
|
45 |
+
top_p=0.95,
|
46 |
+
top_k=40,
|
47 |
+
temperature=0.1,
|
48 |
+
repetition_penalty=1.05,
|
49 |
+
)
|
50 |
+
|
51 |
+
# Giải mã và in kết quả
|
52 |
+
assistant = tokenizer.batch_decode(out_ids[:, input_ids.size(1):], skip_special_tokens=True)[0].strip()
|
53 |
+
print("Assistant: ", assistant)
|