Querybot / app.py
Pyasma's picture
Update app.py
bfe739b
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import gradio as gr
model_name = "TheBloke/Llama-2-7B-Chat-GGML"
use_triton = False
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.to('cuda')
# Creating pipeline
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=512,
temperature=0.6,
top_p=0.95,
repetition_penalty=1.15
)
# Define your system message
system_message = """
You are an AI Query bot that will answer all of the user questions to the best of your ability. You are helpful, respectful, and honest. Your job is to answer the user's query to the best of your ability. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content.
You will be asked questions randomly on any topic. Just answer them to the best of your ability without a single lie and harmful information.
"""
def run(input, history=memory):
return chat.predict(prompt=input)
output = gr.ChatInterface(
run,
inputs="text",
outputs="text",
title="AI Query Bot",
show_system_message=True, # Display the system message
system_message=system_message, # Set the system message
)
output.launch(share=True)