app.py

import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import snapshot_download
from pathlib import Path

def main():
    st.title("Codestral Inference with Hugging Face")
    
    # Download the model files
    st.text("Downloading model...")
    model_id = "mistralai/Codestral-22B-v0.1"
    local_model_path = Path.home().joinpath('mistral_models', model_id)
    local_model_path.mkdir(parents=True, exist_ok=True)

    snapshot_download(repo_id=model_id, allow_patterns=["*.bin", "*.json", "*.model"], local_dir=local_model_path)
    st.success("Model downloaded successfully!")
    
    # Load the model and tokenizer
    st.text("Loading model...")
    tokenizer = AutoTokenizer.from_pretrained(local_model_path)
    model = AutoModelForCausalLM.from_pretrained(local_model_path)
    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
    st.success("Model loaded successfully!")
    
    user_input = st.text_area("Enter your instruction", "Explain Machine Learning to me in a nutshell.")
    max_tokens = st.slider("Max Tokens", min_value=10, max_value=500, value=64)
    temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=0.7)
    
    if st.button("Generate"):
        with st.spinner("Generating response..."):
            result = generate_response(generator, user_input, max_tokens, temperature)
            st.success("Response generated!")
            st.text_area("Generated Response", result, height=200)

def generate_response(generator, user_input, max_tokens, temperature):
    response = generator(user_input, max_new_tokens=max_tokens, do_sample=True, temperature=temperature)
    result = response[0]['generated_text']
    return result

if __name__ == "__main__":
    main()