Shriharsh/llama-3-8B-Instruct-ChatBot

1import gradio as gr2from transformers import AutoModelForCausalLM, AutoTokenizer3import torch4 5model_name = "meta-llama/Meta-Llama-3-8B-Instruct"6device_map = 'auto'7 8def load_model() -> AutoModelForCausalLM:9    return AutoModelForCausalLM.from_pretrained(model_name, device_map=device_map)10 11def load_tokenizer() -> AutoTokenizer:12    return AutoTokenizer.from_pretrained(model_name)13 14def preprocess_messages(message: str, history: list, system_prompt: str) -> dict:15    messages = [{'role': 'system', 'content': system_prompt}, {'role': 'user', 'content': message}]16    prompt = load_tokenizer().apply_chat_template(messages, tokenize=False, add_generation_prompt=True)17    return prompt18 19def generate_text(prompt: str, max_new_tokens: int, temperature: float) -> str:20    model = load_model()21    terminators = [load_tokenizer().eos_token_id, load_tokenizer().convert_tokens_to_ids(['\n'])]22    temp = temperature + 0.123    outputs = model.generate(24        prompt,25        max_new_tokens=max_new_tokens,26        eos_token_id=terminators[0],27        do_sample=True,28        temperature=temp,29        top_p=0.930     )31    return load_tokenizer().decode(outputs[0], skip_special_tokens=True)32 33def chat_function(34    message: str,35    history: list,36    system_prompt: str,37    max_new_tokens: int,38    temperature: float39) -> str:40    prompt = preprocess_messages(message, history, system_prompt)41    return generate_text(prompt, max_new_tokens, temperature)42 43gr.ChatInterface(44    chat_function,45    chatbot=gr.Chatbot(height=400),46    textbox=gr.Textbox(placeholder="Enter message here", container=False, scale=7),47    title="llama-3_8B_Instruct ChatBot",48    description="""Chat with llama-3_8B""",49    theme="soft",50    additional_inputs=[51        gr.Textbox("You shall answer to all the questions as very smart AI", label="System Prompt"),52        gr.Slider(512, 4096, label="Max New Tokens"),53        gr.Slider(0, 1, label="Temperature")54     ]55).launch()