import streamlit as st
from ctransformers import AutoModelForCausalLM
import torch
llm = AutoModelForCausalLM.from_pretrained("TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF", gpu_layers=24)

# Streamlit UI
st.set_page_config(page_title="GenBlog Demo", 
                   page_icon="📚", 
                   layout="centered",
                   initial_sidebar_state='collapsed')
st.header("GenBlog Demo 📚")
st.subheader('This is a demo of the GenBlog assistant', divider='rainbow')
st.write("Enter a blog topic, the number of words and the blog style to generate a blog post.")    

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
st.write(f"Based on the TinyLlama 🦙 model by TheBloke. Device: {device} 🚀")

st.divider()

input_text=st.text_input("Enter the Blog Topic")
col1, col2 = st.columns([5, 5])
with col1:
    no_words=st.text_input("Enter the number of words", value=200)
with col2:
    blog_style=st.selectbox("Select the Blog Style", ["Personal", "Casual", "Proffesional"])


submit=st.button("Generate Blog")
st.divider()
response_field = st.empty()



if submit:
    prompt = f"""
    As a professional writer skilled in {blog_style.lower()} style, your task is to create an engaging and informative blog post about '{input_text}'. Begin with a captivating title that reflects the essence of the topic. Follow with an introduction that sets the stage for the discussion.

    Please ensure the main body delves into the details of the topic, providing insight, analysis, and relevant examples that illuminate the subject for the readers. Conclude with a summary that reinforces the main points and offers a call to action or a thought-provoking question to engage the readers further.

    The post should be well-structured, clearly written, and approximately {no_words} words in total, reflecting a professional tone suitable for an audience interested in {blog_style.lower()} topics.
    """

    st.write_stream(llm(prompt, stream=True, max_new_tokens=256, temperature=0.3))