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()