Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from threading import Thread | |
| from typing import Iterator | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| MAX_MAX_NEW_TOKENS = 1024 | |
| MAX_INPUT_TOKEN_LENGTH = 2048 | |
| # base_model_name = "m-a-p/OpenCodeInterpreter-DS-6.7B" | |
| base_model_name = "m-a-p/OpenCodeInterpreter-DS-1.3B" | |
| model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float32, device_map="cpu", low_cpu_mem_usage=True) | |
| tokenizer = AutoTokenizer.from_pretrained(base_model_name) | |
| def format_prompt(message, history): | |
| system_prompt = "You are OpenCodeInterpreter, you are an expert programmer that helps to write code based on the user request, with concise explanations." | |
| prompt = [] | |
| prompt.append({"role": "system", "content": system_prompt}) | |
| for user_prompt, bot_response in history: | |
| prompt.extend([{"role": "user", "content": user_prompt}, {"role": "assistant", "content": bot_response}]) | |
| prompt.append({"role": "user", "content": message}) | |
| return prompt | |
| def generate(prompt: str, history: list[tuple[str, str]], max_new_tokens: int = 1024, temperature: float = 0.3, | |
| top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1 ) -> Iterator[str]: | |
| temperature = float(temperature) | |
| if temperature < 1e-2: | |
| temperature = 1e-2 | |
| formatted_prompt = [] | |
| formatted_prompt = format_prompt(prompt, history) | |
| input_ids = tokenizer.apply_chat_template(formatted_prompt, return_tensors="pt", add_generation_prompt=True) | |
| if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH: | |
| input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:] | |
| input_ids = input_ids.to(model.device) | |
| streamer = TextIteratorStreamer(tokenizer, timeout=15.0, skip_prompt=True, skip_special_tokens=True) | |
| generation_kwargs = dict({"input_ids": input_ids}, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=False, top_p=top_p, top_k=top_k, | |
| temperature=temperature, num_beams=1, repetition_penalty=repetition_penalty, eos_token_id=tokenizer.eos_token_id) | |
| t = Thread(target=model.generate, kwargs=generation_kwargs ) | |
| t.start() | |
| outputs = [] | |
| for chunk in streamer: | |
| outputs.append(chunk) | |
| yield "".join(outputs).replace("<|EOT|>","") | |
| mychatbot = gr.Chatbot(layout="bubble", avatar_images=["user.png", "botoci.png"], bubble_full_width=False, show_label=False, show_copy_button=True, likeable=True,) | |
| additional_inputs = additional_inputs=[ | |
| gr.Slider( | |
| label="Max new tokens", | |
| minimum=1, | |
| maximum=MAX_MAX_NEW_TOKENS, | |
| step=1, | |
| value=512, | |
| ), | |
| gr.Slider( | |
| label="Temperature", | |
| minimum=0, | |
| maximum=1.0, | |
| step=0.1, | |
| value=0.3, | |
| ), | |
| gr.Slider( | |
| label="Top-p", | |
| minimum=0.05, | |
| maximum=1.0, | |
| step=0.05, | |
| value=0.9, | |
| ), | |
| gr.Slider( | |
| label="Top-k", | |
| minimum=1, | |
| maximum=1000, | |
| step=1, | |
| value=50, | |
| ), | |
| gr.Slider( | |
| label="Repetition penalty", | |
| minimum=1.0, | |
| maximum=2.0, | |
| step=0.05, | |
| value=1, | |
| )] | |
| iface = gr.ChatInterface(fn=generate, | |
| chatbot=mychatbot, | |
| additional_inputs=additional_inputs, | |
| description=" Running on CPU. The response may be slow for cpu environments. 🙏🏻", | |
| retry_btn=None, | |
| undo_btn=None | |
| ) | |
| with gr.Blocks() as demo: | |
| gr.HTML("<center><h1>Tomoniai's Chat with OpenCodeInterpreter</h1></center>") | |
| iface.render() | |
| demo.queue(max_size=10).launch(show_api=False) |