Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,124 Bytes
56ab742 5b9f390 56ab742 5b9f390 56ab742 5b9f390 56ab742 5b9f390 56ab742 5b9f390 77bd93c 5b9f390 77bd93c 5b9f390 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# login as a privileged user.
import os
HF_TOKEN = os.environ.get("HF_TOKEN")
from huggingface_hub import login
login(token=HF_TOKEN)
from threading import Thread
from typing import Iterator
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from pyreft import ReftModel
MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
DESCRIPTION = """\
# ReFT-GOODY-2 on Llama-2 7B Chat
"""
LICENSE = """
<p/>
---
A [GOODY-2](https://www.goody2.ai/chat) imitator built with ReFT, 5 training examples and 30 seconds.
"""
if not torch.cuda.is_available():
DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
if torch.cuda.is_available():
model_id = "meta-llama/Llama-2-7b-chat-hf" # not gated version.
model = AutoModelForCausalLM.from_pretrained(
model_id, device_map="auto", torch_dtype=torch.bfloat16
)
reft_model = ReftModel.load("pyvene/reft_goody2", model, from_huggingface_hub=True)
# a little hacky.
for k, v in reft_model.interventions.items():
v[0].to(model.device)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.use_default_system_prompt = True
prompt_no_input_template = """<s>[INST] <<SYS>>
You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.
<</SYS>>
%s [/INST]
"""
@spaces.GPU
def generate(
message: str,
chat_history: list[tuple[str, str]],
max_new_tokens: int = 1024,
) -> Iterator[str]:
# tokenize and prepare the input
prompt = prompt_no_input_template % message
prompt = tokenizer(prompt, return_tensors="pt").to(model.device)
input_ids = prompt["input_ids"]
attention_mask = prompt["attention_mask"]
if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
attention_mask = attention_mask[:, -MAX_INPUT_TOKEN_LENGTH:]
gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
base_unit_location = input_ids.shape[-1] - 1 # last position
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
generate_kwargs = {
"base": {"input_ids": prompt["input_ids"], "attention_mask": prompt["attention_mask"]},
"unit_locations": {"sources->base": (None, [[[base_unit_location]]])},
"max_new_tokens": max_new_tokens,
"intervene_on_prompt": True,
"streamer": streamer,
"eos_token_id": tokenizer.eos_token_id,
"early_stopping": True,
"do_sample": False,
"repetition_penalty": 1.1,
}
t = Thread(target=reft_model.generate, kwargs=generate_kwargs)
t.start()
outputs = []
for text in streamer:
outputs.append(text)
yield "".join(outputs)
chat_interface = gr.ChatInterface(
fn=generate,
additional_inputs=[
gr.Slider(
label="Max new tokens",
minimum=1,
maximum=MAX_MAX_NEW_TOKENS,
step=1,
value=DEFAULT_MAX_NEW_TOKENS,
)
],
stop_btn=None,
examples=[
["What's 2+2?"],
["Why is the sky blue?"],
["What's Apple's stock price?"],
["Plan a family road trip to Austin"],
],
)
with gr.Blocks(css="style.css") as demo:
gr.Markdown(DESCRIPTION)
gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
chat_interface.render()
gr.Markdown(LICENSE)
if __name__ == "__main__":
demo.queue(max_size=20).launch()
|