import streamlit as st
from llama_cpp import Llama
from huggingface_hub import snapshot_download
import os

# Define the model repository and filename
repo_id = "Divyansh12/check"
filename = "unsloth.F16.gguf"

# Download the model if not already present
model_dir = "model"
if not os.path.exists(model_dir):
    st.write("Downloading the model...")
    snapshot_download(repo_id=repo_id, cache_dir=model_dir)
    st.write("Model downloaded successfully!")

# Since the model file is directly inside the downloaded directory
# Find the exact file path
model_path = os.path.join(model_dir, filename)

# Check if the file exists at the specified location
if not os.path.exists(model_path):
    st.error(f"Model file not found at {model_path}")
else:
    st.write(f"Found model file at {model_path}")

    # Load the GGUF model using llama-cpp
    st.write("Loading the model...")
    llm = Llama.from_pretrained(
        model_path=model_path
    )
    st.write("Model loaded successfully!")

    # Streamlit input for the user to enter a prompt
    prompt = st.text_input("Enter your prompt:", "What is the capital of France?")

    # Generate the response when a prompt is given
    if st.button("Generate Response"):
        with st.spinner("Generating..."):
            response = llm.create_chat_completion(
                messages=[
                    {
                        "role": "user",
                        "content": prompt
                    }
                ]
            )
            # Extract the message content from the response and display it
            st.write("Response:")
            st.write(response['choices'][0]['message']['content'])