Update app.py

app.py

@@ -1,37 +1,30 @@
 import streamlit as st
-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
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
 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")
-
+    # Load the model and tokenizer
+    st.text("Loading model...")
+    tokenizer = AutoTokenizer.from_pretrained("mistralai/Codestral-22B-v0.1")
+    model = AutoModelForCausalLM.from_pretrained("mistralai/Codestral-22B-v0.1")
+    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(user_input, mistral_models_path, max_tokens, temperature)
+            result = generate_response(generator, user_input, max_tokens, temperature)
             st.success("Response generated!")
             st.text_area("Generated Response", result, height=200)
 
-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
-
-    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])
+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()
-