quancute commited on
Commit
c0fd4d8
·
verified ·
1 Parent(s): c031a23

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +103 -0
app.py CHANGED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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='./hf_cache'
18
+ )
19
+
20
+ tokenizer = AutoTokenizer.from_pretrained(MODEL, cache_dir='./hf_cache')
21
+
22
+ # def generate_output(input_text:str,
23
+ # top_p:float=0.95,
24
+ # top_k:int=40,
25
+ # temperature:float=0.1,
26
+ # repetition_penalty:float=1.05,
27
+ # max_new_tokens:int=768):
28
+ # system_prompt = "Bạn là một trợ lí ảo Tiếng Việt về lĩnh vực y tế."
29
+ # conversation = [{"role": "system", "content": system_prompt }]
30
+ # human = f"Vui lòng trả lời câu hỏi sau: {input_text}"
31
+ # conversation.append({"role": "user", "content": human })
32
+
33
+ # # Chuyển các tensor đầu vào sang đúng thiết bị
34
+ # input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(device)
35
+
36
+ # # Tạo đầu ra từ mô hình
37
+ # out_ids = model.generate(
38
+ # input_ids=input_ids,
39
+ # max_new_tokens=768,
40
+ # do_sample=True,
41
+ # top_p=0.95,
42
+ # top_k=40,
43
+ # temperature=0.1,
44
+ # repetition_penalty=1.05,
45
+ # )
46
+
47
+ # # Giải mã và in kết quả
48
+ # assistant = tokenizer.batch_decode(out_ids[:, input_ids.size(1):], skip_special_tokens=True)[0].strip()
49
+ # return assistant
50
+
51
+ def respond(
52
+ message,
53
+ history: list[tuple[str, str]],
54
+ system_message: str,
55
+ max_tokens,
56
+ temperature,
57
+ top_p,
58
+ ):
59
+ sys_prompt = "Bạn là một trợ lí ảo Tiếng Việt về lĩnh vực y tế."
60
+ conversation = [{"role": "system", "content": sys_prompt}]
61
+
62
+
63
+ for val in history:
64
+ if val[0]:
65
+ conversation.append({"role": "user", "content": val[0]})
66
+ if val[1]:
67
+ messages.append({"role": "assistant", "content": val[1]})
68
+
69
+ conversation.append({"role": "user", "content": message})
70
+
71
+ input_ids_list = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(device)
72
+ response = ""
73
+
74
+ for message in tokenizer.batch_decode(model.generate(
75
+ input_ids=input_ids,
76
+ max_new_tokens=max_tokens,
77
+ do_sample=True,
78
+ top_p=top_p,
79
+ temperature=temperature,
80
+ )[:, input_ids_list.size(1):], skip_special_tokens=True):
81
+ token = message.strip()
82
+ response += token
83
+ yield response
84
+
85
+ demo = gr.ChatInterface(
86
+ respond,
87
+ additional_inputs=[
88
+ gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
89
+ gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
90
+ gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
91
+ gr.Slider(
92
+ minimum=0.1,
93
+ maximum=1.0,
94
+ value=0.95,
95
+ step=0.05,
96
+ label="Top-p (nucleus sampling)",
97
+ ),
98
+ ],
99
+ )
100
+
101
+
102
+ if __name__ == "__main__":
103
+ demo.launch()