toxic-llama2 / app.py
freQuensy23's picture
Update app.py
b8d1d3d verified
raw
history blame contribute delete
858 Bytes
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()