Update app.py

app.py

@@ -1,22 +1,37 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_common.protocol.instruct.messages import UserMessage
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
+from mistral_inference.model import Transformer
+from mistral_inference.generate import generate
+from transformers import AutoModelForCausalLM
 
-# Load model and tokenizer
-@st.cache_resource  # Cache the resources to avoid reloading on every run
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("mistralai/Codestral-22B-v0.1")
-    model = AutoModelForCausalLM.from_pretrained("mistralai/Codestral-22B-v0.1")
-    return tokenizer, model
+def main():
+    st.title("Codestral Inference with Hugging Face")
+    
+    mistral_models_path = st.text_input("Enter the path to your Codestral model", "path/to/mistral_models/Codestral-22B-v0.1")
 
-tokenizer, model = load_model()
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+    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(user_input, mistral_models_path, max_tokens, temperature)
+            st.success("Response generated!")
+            st.text_area("Generated Response", result, height=200)
 
-st.title("Chat with Codestral-22B")
-st.write("Ask a question and get a response from the Codestral-22B model.")
+def generate_response(user_input, model_path, max_tokens, temperature):
+    tokenizer = MistralTokenizer.v3()
+    completion_request = ChatCompletionRequest(messages=[UserMessage(content=user_input)])
+    tokens = tokenizer.encode_chat_completion(completion_request).tokens
 
-user_input = st.text_input("You: ", "Type your question here...")
+    model = Transformer.from_folder(model_path)
+    out_tokens, _ = generate([tokens], model, max_tokens=max_tokens, temperature=temperature, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
+
+    result = tokenizer.decode(out_tokens[0])
+    return result
+
+if __name__ == "__main__":
+    main()
 
-if st.button("Send"):
-    with st.spinner("Generating response..."):
-        response = generator(user_input, max_length=100, num_return_sequences=1)
-        st.write("Codestral-22B: " + response[0]["generated_text"])