ChatFriend / app.py
swapnasg's picture
Update app.py
a5cd913 verified
raw
history blame contribute delete
590 Bytes
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
model_id = "tiiuae/falcon-rw-1b"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
def chat_with_user(message, history=[]):
prompt = message
result = generator(prompt, max_new_tokens=200)[0]['generated_text']
history.append((message, result.strip()))
return history, history
gr.ChatInterface(chat_with_user, title="Expert Chat Assistant").launch()