Spaces:
Sleeping
Sleeping
| #!/usr/bin/python3 | |
| # -*- coding: utf-8 -*- | |
| from typing import List, Tuple | |
| from threading import Thread | |
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| from project_settings import project_path | |
| def greet(question: str, history: List[Tuple[str, str]]): | |
| answer = "Hello " + question + "!" | |
| result = history + [(question, answer)] | |
| return result | |
| def chat_with_llm_non_stream(question: str, | |
| history: List[Tuple[str, str]], | |
| pretrained_model_name_or_path: str, | |
| max_new_tokens: int, top_p: float, temperature: float, repetition_penalty: float, | |
| ): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = AutoModelForCausalLM.from_pretrained( | |
| pretrained_model_name_or_path, | |
| trust_remote_code=True, | |
| low_cpu_mem_usage=True, | |
| torch_dtype=torch.bfloat16, | |
| device_map="auto", | |
| offload_folder="./offload", | |
| offload_state_dict=True, | |
| # load_in_4bit=True, | |
| ) | |
| model = model.to(device) | |
| model = model.bfloat16().eval() | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| pretrained_model_name_or_path, | |
| trust_remote_code=True, | |
| # llama不支持fast | |
| use_fast=False if model.config.model_type == "llama" else True, | |
| padding_side="left" | |
| ) | |
| # QWenTokenizer比较特殊, pad_token_id, bos_token_id, eos_token_id 均 为None. eod_id对应的token为<|endoftext|> | |
| if tokenizer.__class__.__name__ == "QWenTokenizer": | |
| tokenizer.pad_token_id = tokenizer.eod_id | |
| tokenizer.bos_token_id = tokenizer.eod_id | |
| tokenizer.eos_token_id = tokenizer.eod_id | |
| input_ids = tokenizer( | |
| question, | |
| return_tensors="pt", | |
| add_special_tokens=False, | |
| ).input_ids.to(device) | |
| bos_token_id = torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long).to(device) | |
| eos_token_id = torch.tensor([[tokenizer.eos_token_id]], dtype=torch.long).to(device) | |
| input_ids = torch.concat([bos_token_id, input_ids, eos_token_id], dim=1) | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| input_ids=input_ids, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=True, | |
| top_p=top_p, | |
| temperature=temperature, | |
| repetition_penalty=repetition_penalty, | |
| eos_token_id=tokenizer.eos_token_id | |
| ) | |
| outputs = outputs.tolist()[0][len(input_ids[0]):] | |
| response = tokenizer.decode(outputs) | |
| response = response.strip().replace(tokenizer.eos_token, "").strip() | |
| return | |
| def main(): | |
| description = """ | |
| chat llm | |
| """ | |
| with gr.Blocks() as blocks: | |
| gr.Markdown(value="gradio demo") | |
| chatbot = gr.Chatbot([], elem_id="chatbot", height=400) | |
| with gr.Row(): | |
| with gr.Column(scale=4): | |
| text_box = gr.Textbox(show_label=False, placeholder="Enter text and press enter", container=False) | |
| with gr.Column(scale=1): | |
| submit_button = gr.Button("💬Submit") | |
| with gr.Column(scale=1): | |
| clear_button = gr.Button( | |
| '🗑️Clear', | |
| variant='secondary', | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| max_new_tokens = gr.Slider(minimum=0, maximum=512, value=512, step=1, label="max_new_tokens"), | |
| with gr.Column(scale=1): | |
| top_p = gr.Slider(minimum=0, maximum=1, value=0.85, step=0.01, label="top_p"), | |
| with gr.Column(scale=1): | |
| temperature = gr.Slider(minimum=0, maximum=1, value=0.35, step=0.01, label="temperature"), | |
| with gr.Column(scale=1): | |
| repetition_penalty = gr.Slider(minimum=0, maximum=2, value=1.2, step=0.01, label="repetition_penalty"), | |
| with gr.Row(): | |
| model_name = gr.Dropdown(choices=["Qwen/Qwen-7B-Chat"], | |
| value="Qwen/Qwen-7B-Chat", | |
| label="model_name", | |
| ) | |
| gr.Examples(examples=["你好"], inputs=text_box) | |
| inputs = [ | |
| text_box, chatbot, model_name, | |
| max_new_tokens, top_p, temperature, repetition_penalty | |
| ] | |
| outputs = [ | |
| chatbot | |
| ] | |
| text_box.submit(chat_with_llm_non_stream, inputs, outputs) | |
| submit_button.click(chat_with_llm_non_stream, inputs, outputs) | |
| clear_button.click( | |
| fn=lambda: ('', ''), | |
| outputs=[text_box, chatbot], | |
| queue=False, | |
| api_name=False, | |
| ) | |
| blocks.queue().launch() | |
| return | |
| if __name__ == '__main__': | |
| main() | |