hf-sme-test / app.py
rawkintrevo's picture
Update app.py
dad6346
import streamlit as st
from transformers import pipeline, BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
import torch
print(torch.cuda.is_available())
tokenizer_name = "tiiuae/falcon-7b-instruct"
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
# Load the Hugging Face model for chatbot
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True)
model = AutoModelForCausalLM.from_pretrained(
"rawkintrevo/hf-sme-falcon-7b",
revision="v0.0.1",
quantization_config=bnb_config,
torch_dtype=torch.float16,
trust_remote_code=True
)
chatbot = pipeline("conversational",
model=model,
tokenizer=tokenizer
)
# Streamlit app title
st.title("Hugging Face Chatbot")
# User input for chat
user_input = st.text_input("You:", "")
if st.button("Ask"):
if user_input:
# Generate a response from the chatbot model
response = chatbot(user_input)[0]['generated_text']
st.text("Chatbot:")
st.write(response)
# Example conversation
st.subheader("Example Conversation:")
st.write("You: Hi, how are you?")
st.write("Chatbot: I'm good, how can I help you today?")