import sys
!{sys.executable} -m pip install streamlit transformers accelerate torch

import streamlit as st

HF_API_TOKEN = "your_huggingface_api_token_here"  # <-- paste your token
MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2"
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
HEADERS = {"Authorization": f"Bearer {HF_API_TOKEN}"}

import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
import torch

st.title("📚 AI Adaptive Learning (Local Small Model)")

MODEL_ID = "microsoft/phi-2"

@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    # Set pad_token_id for the tokenizer if it's not already set, using eos_token_id as a fallback
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id

    # Load the model configuration
    config = AutoConfig.from_pretrained(MODEL_ID)

    # Check if pad_token_id exists in config and set it if not, for Phi models this often needs to be explicitly added
    if not hasattr(config, 'pad_token_id') or config.pad_token_id is None:
        config.pad_token_id = tokenizer.pad_token_id

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        config=config, # Pass the modified config to the model
        torch_dtype=torch.float32,
        device_map="auto"
    )
    return tokenizer, model

tokenizer, model = load_model()

# Input question
user_input = st.text_input("Ask a question:")

if st.button("Submit") and user_input:
    inputs = tokenizer(user_input, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            do_sample=True,
            temperature=0.7
        )

    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)

    st.subheader("AI Answer:")
    st.write(answer)