Spaces:
Running
Running
File size: 2,147 Bytes
8400d16 |
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 |
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import gradio as gr
import torch
MAX_INPUT_TOKEN_LENGTH = 4096
model_id = 'HuggingFaceH4/zephyr-7b-beta'
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map='auto')
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.use_default_system_prompt = False
def generate(input, chat_history=[], system_prompt=False, max_new_tokens=512, temperature=0.5, top_p=0.95, top_k=50, repetition_penalty=1.2):
conversation = []
if system_prompt:
conversation.append({
'role': 'system',
'content': system_prompt
})
for user, assistant in chat_history:
conversation.extend({
'role': 'user',
'content': user
},
{
'role': 'assistant',
'content': assistant
})
conversation.append({
'role': 'user',
'content': input
})
input_ids = tokenizer.apply_chat_template(conversation, return_tensors='pt')
if input_ids.shape[1] > MAXX_INPUT_TOKEN_LENGTH:
input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
gr.Warning(f"Trimed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
input_ids = input_ids.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,
top_p=top_p,
top_k=top_k,
temperature=temperature,
num_beams=1,
repetition_penalty=repetition_penalty
)
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
outputs = []
for text in streamer:
outputs.append(text)
yield ''.join(outputs)
chat_interface = gr.ChatInterface(
fn=generate,
examples=[
'What is GPT?',
'What is Life?',
'Who is Alan Turing'
]
)
chat_interface.queue(max_size=20).launch() |