import streamlit as st
from transformers import pipeline

# generator = pipeline(model="meta-llama/Llama-2-7b-chat-hf") # gated model
# generator = pipeline(model="TheBloke/CapybaraHermes-2.5-Mistral-7B-GPTQ") # Error
generator = pipeline(model="xai-org/grok-1")
prompt = st.text_area('Your prompt: ')
res = generator(prompt)

if res:
    out = pipe(res)
    st.json(out)