Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import pipeline | |
| import torch | |
| def main(): | |
| st.title("Chatbot with Hugging Face Model") | |
| # If not saved, load the model and save it | |
| pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta", torch_dtype=torch.bfloat16, device_map="auto") | |
| user_input = st.text_area("Enter Text") | |
| # Define chat messages | |
| messages = [ | |
| {"role": "system", "content": "You are a friendly chatbot who always responds in kind way."}, | |
| {"role": "user", "content": user_input}, | |
| ] | |
| submit = st.button('Generate') | |
| if submit: | |
| # Generate response | |
| prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95) | |
| # Display generated text | |
| st.text(outputs[0]["generated_text"]) | |
| if __name__ == "__main__": | |
| main() |