Spaces:
Sleeping
Sleeping
File size: 2,762 Bytes
d9a9459 613da68 d9a9459 d41a141 d9a9459 d41a141 d9a9459 d41a141 d9a9459 |
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 52 53 54 55 56 57 58 59 60 61 62 |
import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
st.set_page_config(page_title="Simple LLM Chatbot", page_icon="icon.png")
def load_model_tokenizer(model_name, hf_api_key):
if model_name == "LLaMa-2B":
model_name="llmware/bling-sheared-llama-2.7b-0.1"
model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_api_key)
tokenizer = AutoTokenizer.from_pretrained(model_name, tokenizer=hf_api_key)
elif model_name == "Red-Pajamas-3b":
model_name = "llmware/bling-red-pajamas-3b-0.1"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
return (model,tokenizer)
def generate_response(prompt_input, model, tokenizer):
inputs = tokenizer.encode_plus(prompt_input, return_tensors="pt")
# Generate the response from the model with additional parameters
outputs = model.generate(**inputs, max_length=max_length, do_sample=True ,temperature=temperature)
response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
return response
st.set_page_config(page_title="Learn Geoscience")
with st.sidebar:
st.title('Learn Geoscience Chat')
if 'hf_key' in st.secrets:
st.success('Huggingface API key provided', icon='✅')
hf_api_key = st.secrets['hf_key']
else:
hf_api_key = st.text_input('Enter Huggingface API Key:', type='password')
if not hf_api_key:
st.warning('Please enter Huggingface API key!', icon='⚠️')
else:
st.success('Proceed to entering your prompt message!', icon='👉')
max_length = st.slider("Max Length", 10, 100, 50)
temperature = st.slider("Temperature", 0.0, 1.0, 0.7)
if "messages" not in st.session_state.keys():
st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
model_name = st.radio("Select model to chat", options=["LLaMa-2B", "Red-Pajamas-3b"], horizontal=True, key='model_selection')
model, tokenizer = load_model_tokenizer(model_name, hf_api_key)
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
if prompt := st.chat_input(disabled = not hf_api_key):
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("user"):
st.write(prompt)
if st.session_state.messages[-1]["role"] != "assistant":
with st.chat_message("assistant"):
with st.spinner("Thinking..."):
response = generate_response(prompt, model, tokenizer)
st.write(response)
message = {"role": "assistant", "content": response}
st.session_state.messages.append(message)
|