import os import json import torch import spaces import gradio as gr from threading import Thread from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer from huggingface_hub import HfApi from datetime import datetime MODEL_ID = os.environ.get("MODEL_ID") DATASET_REPO = os.environ.get("DATASET_REPO") DESCRIPTION = os.environ.get("DESCRIPTION") PROMPT = os.environ.get("PROMPT") ON_LOAD=""" async()=>{ alert("Before using the service, users must agree to the following terms:\\n\\nPlease note that the model presented here is an experimental tool that is still being developed and improved.\\n\\nMeasures have been taken during the model creation process to minimizing the risk of generating vulgar, prohibited or inappropriate content. However, in rare cases, unwanted content may be generated. If you encounter any content that is deemed inappropriate or violates our policies, please contact us to report it. Your information will enable us to take further steps to improve and develop the model to make it safe and user-friendly.\\n\\nYou must not use the model for illegal, harmful, violent, racist or sexual purposes. Please do not send any private information. The website collects user dialogue data and reserves the right to distribute it under the Creative Commons Attribution (CC-BY) or similar license."); } """ api = HfApi() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') @spaces.GPU() def generate(instruction, stop_tokens, repetition_penalty, max_new_tokens): streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) input_ids = tokenizer.apply_chat_template(instruction, return_tensors="pt", add_generation_prompt=True,).to(device) if input_ids.shape[1] > 4096: input_ids = input_ids[:, -4096:] generate_kwargs = dict( input_ids = input_ids, streamer=streamer, do_sample=False, max_new_tokens=max_new_tokens, repetition_penalty=repetition_penalty, ) t = Thread(target=model.generate, kwargs=generate_kwargs) t.start() outputs = [] for new_token in streamer: if new_token in stop_tokens: break outputs.append(new_token.replace("<|im_end|>", "")) yield "".join(outputs) def predict(message, history): system_prompt = PROMPT repetition_penalty = 1.1 max_new_tokens = 1024 stop_tokens = ["<|endoftext|>", "<|im_end|>"] conversation = [] conversation.append( { "role": "system", "content": system_prompt, } ) for user, assistant in history: conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}]) conversation.append({"role": "user", "content": message}) print(conversation) for output_text in generate(conversation, stop_tokens, repetition_penalty, max_new_tokens): if output_text in stop_tokens: break yield output_text hfapi = HfApi() day=datetime.now().strftime("%Y-%m-%d") timestamp=datetime.now().timestamp() dd={ 'message': message, 'history': history, 'system_prompt':system_prompt, 'max_new_tokens':max_new_tokens, 'repetition_penalty':repetition_penalty, 'instruction':conversation, 'output':output_text, 'precision': 'auto '+str(model.dtype), } hfapi.upload_file( path_or_fileobj=json.dumps(dd, indent=2, ensure_ascii=False).encode('utf-8'), path_in_repo=f"{day}/{timestamp}.json", repo_id=DATASET_REPO, repo_type="dataset", commit_message=f"X", run_as_future=True ) def vote(chatbot, data: gr.LikeData): day=datetime.now().strftime("%Y-%m-%d") timestamp=datetime.now().timestamp() api.upload_file( path_or_fileobj=json.dumps({"history":chatbot, 'index': data.index, 'liked': data.liked}, indent=2, ensure_ascii=False).encode('utf-8'), path_in_repo=f"liked/{day}/{timestamp}.json", repo_id=DATASET_REPO, repo_type="dataset", commit_message=f"L", run_as_future=True ) with gr.Blocks(js=ON_LOAD) as demo: chatbot = gr.Chatbot(label="Chatbot", likeable=True, render=False) chatbot.like(vote, [chatbot], None) gr.ChatInterface( predict, chatbot=chatbot, title="MKLLM-7B-Instruct", description=DESCRIPTION, examples=[ ["Кој си ти?"], ["Колку е 3+6/3-1?"], ["Како би го населиле Марс? Биди краток."], ["Напиши ми pyhon функција која прима низа броеви и проверува кои броеви се деливи со 7, оние што се деливи ги сместуваме во нова листа која ја враќаме назад."] ] ) if __name__ == "__main__": tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) tokenizer.eos_token = "<|im_end|>" tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map=device, torch_dtype='auto', ) demo.queue(max_size=50).launch()