File size: 3,858 Bytes
d79aeef fb62856 d79aeef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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) |