File size: 858 Bytes
ad1e673
3c4ad3a
 
b8d1d3d
ad1e673
2edad32
b8d1d3d
2219bdc
b8d1d3d
ad1e673
3c4ad3a
 
 
 
 
 
 
 
ad1e673
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import gradio as gr
import transformers
import peft
import os

model_id = 'freQuensy23/toxic-llama2'
model = peft.AutoPeftModelForCausalLM.from_pretrained(model_id, token=os.getenv('hf_token'))
model.to_bettertransformer()
tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, token=os.getenv('hf_token'))


def generate(text, temp):
    input_ids = tokenizer(f"User: {text}\nBot:", return_tensors='pt').input_ids
    generated_ids = model.generate(input_ids=input_ids.to(model.device), temperature=temp, max_new_tokens=64)[0][len(input_ids[0]):]
    return tokenizer.decode(generated_ids).split('\n')[0]

iface = gr.Interface(concurrency_limit=2, fn=generate, inputs=[gr.Textbox(lines=5, placeholder="Type your prompt here...", value='''I am clever?'''), gr.Slider(0.1, 1.5, value=1.1)], 
                     outputs=gr.Textbox())
iface.launch()