test-deployment / src /streamlit_app.py
ropoir's picture
mini chatbot
71b93f9
raw
history blame contribute delete
512 Bytes
import streamlit as st
from transformers import pipeline
st.title("Mini Chatbot")
@st.cache_resource
def load_model():
pipe = pipeline(
"text-generation",
model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
device_map="cpu"
)
return pipe
pipe = load_model()
prompt = st.text_input("Your message")
if st.button("Send"):
output = pipe(
prompt,
max_new_tokens=100,
do_sample=True,
temperature=0.7
)
st.write(output[0]["generated_text"])