import gradio as gr from huggingface_hub import InferenceClient import random models=[ "google/gemma-7b", "google/gemma-7b-it", "google/gemma-2b", "google/gemma-2b-it" ] clients=[ InferenceClient(models[0]), InferenceClient(models[1]), InferenceClient(models[2]), InferenceClient(models[3]), ] def format_prompt(message, history): prompt = "" if history: for user_prompt, bot_response in history: prompt += f"[INST] {user_prompt} [/INST]" prompt += f" {bot_response} " prompt += f"[INST] {message} [/INST]" return prompt def chat_inf(system_prompt,prompt,history,client_choice): #token max=8192 client=clients[int(client_choice)-1] if not history: history = [] hist_len=0 if history: hist_len=len(history) print(hist_len) seed = random.randint(1,1111111111111111) generate_kwargs = dict( temperature=0.9, max_new_tokens=6000, top_p=0.95, repetition_penalty=1.0, do_sample=True, seed=seed, ) formatted_prompt=prompt #formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history) stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False) output = "" for response in stream: output += response.token.text yield [(prompt,output)] history.append((prompt,output)) yield history with gr.Blocks() as app: with gr.Group(): chat_b = gr.Chatbot() with gr.Row(): with gr.Column(scale=3): inp = gr.Textbox(label="Prompt") sys_inp = gr.Textbox(label="System Prompt (optional)") btn = gr.Button("Chat") with gr.Column(scale=1): with gr.Group(): stop_btn=gr.Button("Stop") clear_btn=gr.Button("Clear") client_choice=gr.Dropdown(label="Models",type='index',choices=[c for c in models],value=models[0],interactive=True) go=btn.click(chat_inf,[sys_inp,inp,chat_b,client_choice],chat_b) stop_btn.click(None,None,None,cancels=go) app.launch()