mikeee's picture
Duplicate from FredZhang7/chat-test
360d9e4
raw history blame
No virus
1.34 kB
import os
os.system("pip install --upgrade torch transformers sentencepiece scipy cpm_kernels accelerate bitsandbytes")
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
# tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
# model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("openchat/openchat_v2_w")
model = AutoModelForCausalLM.from_pretrained("openchat/openchat_v2_w", load_in_8bit_fp32_cpu_offload=True, load_in_8bit=True)
model.float()
model = model.eval()
model_path = model.config._dict['model_name_or_path']
model_size_gb = os.path.getsize(model_path) / (1024 * 1024 * 1024)
print(f"The model '{model_name}' is taking approximately {model_size_gb:.2f} GB of disk space.")
# with gr.Blocks() as demo:
# chatbot = gr.Chatbot()
# msg = gr.Textbox()
# clear = gr.ClearButton([msg, chatbot])
# def respond(message, chat_history):
# response, chat_history = model.chat(tokenizer, message, history=chat_history, temperature=0.7, repetition_penalty=1.2, max_length=128)
# chat_history.append((message, response))
# return "", chat_history
# msg.submit(respond, [msg, chatbot], [msg, chatbot])
# if __name__ == "__main__":
# demo.launch()