Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 
-# Cache model loading
+# Function to load models with caching
 @st.cache_resource
 def load_model(model_name):
     try:
@@ -16,8 +16,8 @@ def load_model(model_name):
 # Primary model (DeepSeek-R1)
 MODEL_NAME = "deepseek-ai/DeepSeek-R1"
 
-# Fallback model (Mistral-7B-Instruct, optimized for CPU)
-FALLBACK_MODEL = "mistralai/Mistral-7B-Instruct"
+# Fallback model (Mistral-7B-v0.1, smaller and CPU-friendly)
+FALLBACK_MODEL = "mistralai/Mistral-7B-v0.1"
 
 st.title("Text-to-Text AI with DeepSeek-R1")
 
@@ -41,3 +41,4 @@ if st.button("Generate Response"):
             st.error("No working model available.")
     else:
         st.warning("Please enter a prompt.")
+