Update app.py

app.py

@@ -1,14 +1,13 @@
 import gradio as gr
-from transformers import AutoTokenizer
-from mixtrl import MixtralConfig, MixtralForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 from peft import PeftModel
 
 # Load the tokenizer and model
 model_name = "JuliaUpton/Math_AI"
 model_architecture = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 mixtral_tokenizer = AutoTokenizer.from_pretrained(model_architecture)
-mixtral_config = MixtralConfig.from_pretrained(model_architecture)
-merged_model = PeftModel.from_pretrained(model_name, model=MixtralForCausalLM.from_pretrained(model_architecture, config=mixtral_config, device_map="auto"))
+mixtral_config = AutoConfig.from_pretrained(model_architecture, trust_remote_code=True)
+merged_model = PeftModel.from_pretrained(model_name, model=AutoModelForCausalLM.from_pretrained(model_architecture, config=mixtral_config, trust_remote_code=True, device_map="auto"))
 
 # Define the input formatting function
 def input_from_text(instruction):