georgesung's picture
Update app.py
cb96876
from transformers import LlamaForCausalLM, LlamaTokenizer, pipeline
import torch
import gradio as gr
# LLM helper functions
def get_response_text(data):
text = data[0]["generated_text"]
assistant_text_index = text.rfind('### RESPONSE:')
if assistant_text_index != -1:
text = text[assistant_text_index+len('### RESPONSE:'):].strip()
return text
def get_llm_response(prompt, pipe):
raw_output = pipe(prompt)
text = get_response_text(raw_output)
return text
# Load LLM
model_id = "georgesung/llama2_7b_chat_uncensored"
tokenizer = LlamaTokenizer.from_pretrained(model_id)
model = LlamaForCausalLM.from_pretrained(model_id, device_map="auto", load_in_8bit=True)
# Llama tokenizer missing pad token
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_length=4096, # Llama-2 default context window
temperature=0.7,
top_p=0.95,
repetition_penalty=1.15
)
with gr.Blocks() as demo:
gr.Markdown("""
# Chat with llama2_7b_chat_uncensored
NOTICE: I will pause this space on Monday, July 24, around noon UTC. Since it costs $$ to run :)
If you wish to run this space yourself, you can duplicate this space and run it on a T4 small instance.
""")
chatbot = gr.Chatbot()
msg = gr.Textbox()
clear = gr.Button("Clear")
def hist_to_prompt(history):
prompt = ""
for human_text, bot_text in history:
prompt += f"### HUMAN:\n{human_text}\n\n### RESPONSE:\n"
if bot_text:
prompt += f"{bot_text}\n\n"
return prompt
def get_bot_response(text):
bot_text_index = text.rfind('### RESPONSE:')
if bot_text_index != -1:
text = text[bot_text_index + len('### RESPONSE:'):].strip()
return text
def user(user_message, history):
return "", history + [[user_message, None]]
def bot(history):
#bot_message = random.choice(["How are you?", "I love you", "I'm very hungry"])
#history[-1][1] = bot_message + '</s>'
hist_text = hist_to_prompt(history)
print(hist_text)
bot_message = get_llm_response(hist_text, pipe) + tokenizer.eos_token
history[-1][1] = bot_message # add bot message to overall history
return history
msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
bot, chatbot, chatbot
)
clear.click(lambda: None, None, chatbot, queue=False)
demo.queue()
demo.launch()