Spaces:
Sleeping
Sleeping
File size: 1,612 Bytes
edc19d9 5bdeca3 edc19d9 e7c435f edc19d9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import torch
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
# Model name
model_name = "ybelkada/falcon-7b-sharded-bf16"
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
# Load model in CPU mode
device = "cpu" # Hugging Face Spaces does not provide free GPUs
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16, # Use float16 for lower memory usage
device_map=device
)
# Streamlit UI
st.title("🦜 Falcon-7B Chatbot")
st.write("Ask me anything!")
# Store chat history
if "chat_history" not in st.session_state:
st.session_state.chat_history = []
# User input
user_input = st.text_input("You:", "")
if user_input:
# Tokenize input
inputs = tokenizer(user_input, return_tensors="pt")
inputs.pop("token_type_ids", None) # Remove token_type_ids to avoid errors
inputs = {key: value.to(device) for key, value in inputs.items()} # Move inputs to device
# Generate response
with torch.no_grad():
output = model.generate(**inputs, max_length=200, do_sample=True, top_k=50, top_p=0.95)
# Decode response
response = tokenizer.decode(output[:, inputs["input_ids"].shape[-1]:][0], skip_special_tokens=True)
# Store and display chat history
st.session_state.chat_history.append(("You", user_input))
st.session_state.chat_history.append(("Bot", response))
# Display chat history
for sender, message in st.session_state.chat_history:
st.write(f"**{sender}:** {message}")
|