|
import gradio as gr |
|
import os |
|
import spaces |
|
from transformers import GemmaTokenizer, AutoModelForCausalLM |
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer |
|
from threading import Thread |
|
|
|
|
|
HF_TOKEN = os.environ.get("HF_TOKEN", None) |
|
|
|
|
|
DESCRIPTION = ''' |
|
<div> |
|
<h1 style="text-align: center;">Test Model</h1> |
|
</div> |
|
''' |
|
|
|
LICENSE = """ |
|
<p/> |
|
|
|
--- |
|
""" |
|
|
|
PLACEHOLDER = """ |
|
""" |
|
|
|
|
|
css = """ |
|
h1 { |
|
text-align: center; |
|
display: block; |
|
} |
|
|
|
#duplicate-button { |
|
margin: auto; |
|
color: white; |
|
background: #1565c0; |
|
border-radius: 100vh; |
|
} |
|
""" |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("Orenguteng/Llama-3-8B-Lexi-Uncensored") |
|
model = AutoModelForCausalLM.from_pretrained("Orenguteng/Llama-3-8B-Lexi-Uncensored", device_map="auto") |
|
terminators = [ |
|
tokenizer.eos_token_id, |
|
tokenizer.convert_tokens_to_ids("<|eot_id|>") |
|
] |
|
|
|
@spaces.GPU(duration=120) |
|
def chat_llama3_8b(message: str, |
|
history: list, |
|
temperature: float, |
|
max_new_tokens: int, |
|
top_p: float, |
|
system_prompt: str |
|
) -> str: |
|
""" |
|
Generate a streaming response using the llama3-8b model. |
|
Args: |
|
message (str): The input message. |
|
history (list): The conversation history used by ChatInterface. |
|
temperature (float): The temperature for generating the response. |
|
max_new_tokens (int): The maximum number of new tokens to generate. |
|
top_p (float): The top_p value for nucleus sampling. |
|
system_prompt (str): The system prompt to guide the conversation. |
|
Returns: |
|
str: The generated response. |
|
""" |
|
conversation = [{"role": "system", "content": system_prompt}] |
|
for user, assistant in history: |
|
conversation.append({"role": "user", "content": user}) |
|
conversation.append({"role": "assistant", "content": assistant}) |
|
|
|
conversation.append({"role": "user", "content": message}) |
|
|
|
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device) |
|
attention_mask = input_ids.ne(tokenizer.pad_token_id).long() |
|
|
|
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True) |
|
|
|
generate_kwargs = dict( |
|
input_ids= input_ids, |
|
attention_mask=attention_mask, |
|
streamer=streamer, |
|
max_new_tokens=max_new_tokens, |
|
do_sample=True, |
|
temperature=temperature, |
|
top_p=top_p, |
|
eos_token_id=terminators, |
|
pad_token_id=tokenizer.eos_token_id, |
|
) |
|
|
|
if temperature == 0: |
|
generate_kwargs['do_sample'] = False |
|
|
|
t = Thread(target=model.generate, kwargs=generate_kwargs) |
|
t.start() |
|
|
|
outputs = [] |
|
first_chunk = True |
|
for text in streamer: |
|
if first_chunk and text.startswith("assistant"): |
|
text = text[len("assistant"):].lstrip(": \n") |
|
first_chunk = False |
|
|
|
outputs.append(text) |
|
yield "".join(outputs) |
|
|
|
|
|
chatbot=gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface', type='messages') |
|
|
|
with gr.Blocks(fill_height=True, css=css) as aida: |
|
|
|
gr.Markdown(DESCRIPTION) |
|
gr.ChatInterface( |
|
fn=chat_llama3_8b, |
|
chatbot=None, |
|
fill_height=True, |
|
additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False, render=False), |
|
additional_inputs=[ |
|
gr.Slider(minimum=0, |
|
maximum=1, |
|
step=0.1, |
|
value=0.8, |
|
label="Temperature", |
|
render=False), |
|
gr.Slider(minimum=128, |
|
maximum=4096, |
|
step=1, |
|
value=4096, |
|
label="Max new tokens", |
|
render=False ), |
|
gr.Slider(minimum=0, |
|
maximum=1, |
|
step=0.1, |
|
value=0.9, |
|
label="Top_p", |
|
render=False), |
|
gr.Textbox(lines=2, |
|
placeholder="Enter system prompt here...", |
|
label="System Prompt", |
|
render=False), |
|
], |
|
examples=[ |
|
['Who Are you?'] |
|
|
|
], |
|
cache_examples=False, |
|
) |
|
|
|
if __name__ == "__main__": |
|
aida.launch(share=True) |
|
|