Update app.py

app.py

@@ -1,23 +1,26 @@
-# Ensure you don't load bitsandbytes if you are on CPU
-from transformers import AutoTokenizer
-from unsloth import FastLanguageModel
 import torch
+from unsloth import FastLanguageModel
 
+# Load the model on the CPU
 model_name = "AIModels24/Indian_Constitution"  # Replace with your actual model path
 
-@st.cache_resource
-def load_model():
-    # Load the tokenizer and model with CPU optimizations
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    
-    # Skip the 'bitsandbytes' specific configurations for CPU use
-    model = FastLanguageModel.from_pretrained(
-        model_name=model_name,
-        max_seq_length=2048,
-        load_in_4bit=False,  # Disable 4-bit quantization (you can adjust this for CPU)
-        dtype=torch.float32,  # Use float32 (default for CPU)
-    )
-    
-    return model, tokenizer
+# Explicitly set the device to CPU
+device = torch.device('cpu')
+
+# Load model without GPU dependencies
+model = FastLanguageModel.from_pretrained(
+    model_name=model_name,
+    max_seq_length=2048,
+    load_in_4bit=False,  # Disable 4-bit quantization (required for CPU)
+    dtype=torch.float32,  # Use float32 (default for CPU)
+)
+
+model = model.to(device)  # Ensure the model is loaded to CPU
 
-model, tokenizer = load_model()
+# Example inference function
+def generate_text(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = inputs.to(device)  # Move inputs to CPU
+    with torch.no_grad():
+        outputs = model.generate(inputs['input_ids'], max_length=200)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)