chatbot-v03 / app.py
GiteshUjgaonkar's picture
Update app.py
401f91c verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from peft import PeftModel
# Load tokenizer and base model
base_model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-rw-1b", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-rw-1b")
# Load adapter
model = PeftModel.from_pretrained(base_model, "GiteshUjgaonkar/chatbot-v02")
# Pipeline
generator = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_length=300,
temperature=0.7,
top_p=0.9,
do_sample=True,
)
def chat(user_input, history=[]):
system_prompt = "You are a kind, helpful, and friendly AI assistant."
full_input = f"<|system|>\n{system_prompt}\n<|user|>\n{user_input}\n<|assistant|>\n"
output = generator(full_input)[0]["generated_text"]
response = output.split("<|assistant|>\n")[-1].strip()
return response
gr.Interface(fn=chat, inputs="text", outputs="text", title="Gitesh's Chatbot").launch()