fallback to mistral

models/mistral_model.py

@@ -24,22 +24,23 @@ class MistralModel:
     
     def _initialize_model(self):
         """Initialize Mistral model with optimizations"""
-        print("Loading TinyGPT2Model model...")
+        print("Loading Mistral model...")
         
-        model_id = "sshleifer/tiny-gpt2"
+        model_id = "mistralai/Mistral-7B-Instruct-v0.2"
         
         # Load tokenizer
-        MistralModel._tokenizer = AutoTokenizer.from_pretrained(model_id, token=HUGGINGFACE_TOKEN,use_fast=False)
+        MistralModel._tokenizer = AutoTokenizer.from_pretrained(model_id, token=HUGGINGFACE_TOKEN)
         
         # Load model with optimizations
         MistralModel._model = AutoModelForCausalLM.from_pretrained(
             model_id,
             token=HUGGINGFACE_TOKEN,
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float16,
+            device_map="auto",
             load_in_8bit=True  # Use 8-bit quantization for memory efficiency
         )
         
-        print("TinyGPT2Model loaded successfully!")
+        print("Mistral model loaded successfully!")
     
     def generate(
         self,