import streamlit as st
from mlx_lm import load, generate
from huggingface_hub import login
import os

# os.environ["HF_HOME"] = "your_huggingface_token"

token = os.getenv("HF_TOKEN")

login(token=token)

# Load your model and tokenizer


# response = generate(model, tokenizer, prompt=prompt, verbose=True)

st.title("Your F1 Asli")

# User input
user_input = st.text_input("Enter text:")


if st.button("Submit"):
    print("Habibi mereko dhundti kya")
    st.write("mai print toh ho hi raha hu na bantai")

    # Tokenize input and make predictions
    # inputs = tokenizer(user_input, return_tensors="pt")
    # outputs = model(**inputs)
    model, tokenizer = load("Rafii/f1llama")

    prompt="hello"

    if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None:
        messages = [{"role": "user", "content": prompt}]
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        response = generate(model, tokenizer, prompt=user_input, verbose=True)

    st.write(response)