Update app.py

app.py

@@ -1,8 +1,9 @@
-from transformers import  AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
+from transformers import  AutoConfig, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, MistralForCausalLM
 from peft import PeftModel, PeftConfig
 import torch
 import gradio as gr
-
+  
+  
 # Use the base model's ID
 base_model_id = "mistralai/Mistral-7B-v0.1"
 model_directory = "Tonic/mistralmed"
@@ -21,7 +22,7 @@ tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", trust_rem
 
 # Load the PEFT model
 peft_config = PeftConfig.from_pretrained("Tonic/mistralmed", token="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")
-peft_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", trust_remote_code=True)
+peft_model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1", trust_remote_code=True)
 peft_model = PeftModel.from_pretrained(base_model, "Tonic/mistralmed", token="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")
 
 class ChatBot: