from dotenv import load_dotenv
import os
import streamlit as st
from transformers import pipeline

# Load environment variables from .env file
load_dotenv()

# Get the Hugging Face API token from environment variables
hf_token = os.getenv("textgen")

if not hf_token:
    st.error("Hugging Face API token is not set. Please set the HUGGINGFACE_HUB_TOKEN environment variable.")
else:
    # Initialize the Hugging Face pipeline with authentication
    pipe = pipeline("text-generation", model="mistralai/mathstral-7B-v0.1", use_auth_token=hf_token)

    # Function to get response from the model
    def get_response(input_text, keywords, blog_style, max_new_tokens=250):
        # Prompt Template
        template = """
        Generate technical project ideas for {blog_style} job profile for a topic {input_text} using these keywords: {keywords}.
        """
        
        prompt = template.format(blog_style=blog_style, input_text=input_text, keywords=keywords)
        
        # Generate the response from the model
        response = pipe(prompt, max_new_tokens=max_new_tokens)
        return response[0]['generated_text']  # Extract the generated text

    # Streamlit configuration
    st.set_page_config(page_title="Generate Project Idea",
                       page_icon='🤖',
                       layout='centered',
                       initial_sidebar_state='collapsed')

    st.header("Generate Project Idea 🤖")

    input_text = st.text_input("Enter the Topic")

    # Creating two more columns for additional fields
    col1, col2 = st.columns([5, 5])

    with col1:
        keywords = st.text_input('Keywords')
    with col2:
        blog_style = st.selectbox('Generating project idea for',
                                  ('Researchers', 'Data Scientist', 'Software Developer', 'Common People'), index=0)

    submit = st.button("Generate")

    # Final response
    if submit:
        response = get_response(input_text, keywords, blog_style)
        st.write(response)