import streamlit as st
from transformers import pipeline

st.set_page_config(
    page_title="LLAMA2",
    page_icon="🦙",
    layout="centered",
    initial_sidebar_state="expanded"
)

st.markdown(
    body = "# <center>FINE-TUNED LLAMA2 7B CHAT HF</center>",
    unsafe_allow_html = True
)
text="""4Bit QLoRA Fine-Tuned LLM on English Quotes Dataset."""

st.caption(text)

model_id = "iamsubrata/Llama-2-7b-chat-hf-sharded-bf16-fine-tuned"

LLM = pipeline(
    task = "text-generation",
    model = model_id
)

user_input = st.text_area(
    label = "Prompt...",
    value = "Tell me a joke."
)

button  = st.button(
    label = "Generate",
    key = "generate"
)

if user_input and button:
    st.write(LLM(str(user_input)))