Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import pipeline | |
| st.title("Mini Chatbot") | |
| 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"]) |