thviet79 commited on
Commit
421e099
·
verified ·
1 Parent(s): 6088404

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +53 -0
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)