|
import gradio as gr |
|
import os |
|
import spaces |
|
from transformers import GemmaTokenizer, AutoModelForCausalLM |
|
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer |
|
from threading import Thread |
|
import torch |
|
|
|
|
|
HF_TOKEN = os.environ.get("HF_TOKEN", None) |
|
|
|
|
|
DESCRIPTION = ''' |
|
<div> |
|
<h1 style="text-align: center;">Loki ๐๏ธ</h1> |
|
<p>This uses Llama 3 and GPT-4o as generation, both of these make the final generation. <a href="https://huggingface.co/meta-llama/Meta-Llama-3-8B"><b>Llama3-8b</b></a>and <a href="https://platform.openai.com/docs/models/gpt-4o"><b>GPT-4o</b></a></p> |
|
</div> |
|
''' |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct") |
|
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct", torch_dtype=torch.float16).to('cuda') |
|
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 |
|
) -> 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. |
|
Returns: |
|
str: The generated response. |
|
""" |
|
conversation = [] |
|
for user, assistant in history: |
|
conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}]) |
|
conversation.append({"role": "user", "content": message}) |
|
|
|
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device) |
|
|
|
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True) |
|
|
|
generate_kwargs = dict( |
|
input_ids= input_ids, |
|
streamer=streamer, |
|
max_new_tokens=max_new_tokens, |
|
do_sample=True, |
|
temperature=temperature, |
|
eos_token_id=terminators, |
|
) |
|
|
|
if temperature == 0: |
|
generate_kwargs['do_sample'] = False |
|
|
|
t = Thread(target=model.generate, kwargs=generate_kwargs) |
|
t.start() |
|
|
|
outputs = [] |
|
for text in streamer: |
|
outputs.append(text) |
|
return outputs |
|
|
|
|
|
|
|
chatbot=gr.Chatbot(height=600, label='Loki AI') |
|
|
|
with gr.Blocks(fill_height=True) as demo: |
|
|
|
gr.Markdown(DESCRIPTION) |
|
gr.ChatInterface( |
|
fn=chat_llama3_8b, |
|
chatbot=chatbot, |
|
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.95, |
|
label="Temperature", |
|
render=False), |
|
gr.Slider(minimum=128, |
|
maximum=4096, |
|
step=1, |
|
value=512, |
|
label="Max new tokens", |
|
render=False ), |
|
], |
|
examples=[ |
|
['How to setup a human base on Mars? Give short answer.'], |
|
['Explain theory of relativity to me like Iโm 8 years old.'], |
|
['What is 9,000 * 9,000?'], |
|
['Write a pun-filled happy birthday message to my friend Alex.'], |
|
['Justify why a penguin might make a good king of the jungle.'] |
|
], |
|
cache_examples=False, |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.launch() |