Update app.py

app.py

@@ -1,31 +1,41 @@
 import gradio as gr
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig
 import torch
 
-# Check if GPU is available
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+# Set device to CPU since GPU quantization is unavailable
+device = 'cpu'
 
-# Load the model with 8-bit precision
+# Set up 8-bit quantization with BitsAndBytesConfig
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,  # Enable 8-bit quantization
+    llm_int8_enable_fp32_cpu_offload=True  # Use CPU for 8-bit quantization operations
+)
+
+# Load the model with quantization configuration
 model_name = "Rahmat82/DistilBERT-finetuned-on-emotion"
-model = AutoModelForSequenceClassification.from_pretrained(model_name, load_in_8bit=True, device_map="auto")
+model = AutoModelForSequenceClassification.from_pretrained(
+    model_name,
+    quantization_config=quantization_config,
+    device_map={"": device}  # Ensures everything runs on CPU
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 
 def predict(query: str) -> dict:
     inputs = tokenizer(query, return_tensors='pt')
-    inputs = {k: v.to(device) for k, v in inputs.items()}  # Move inputs to the appropriate device
+    inputs = {k: v.to(device) for k, v in inputs.items()}  # Ensure inputs are on CPU
     outputs = model(**inputs)
     outputs = torch.sigmoid(outputs.logits)
     outputs = outputs.detach().cpu().numpy()
     
     # Define label to ID mapping
     label2ids = {
-                "sadness": 0,
-                "joy": 1,
-                "love": 2,
-                "anger": 3,
-                "fear": 4,
-                "surprise": 5,
-                }
+        "sadness": 0,
+        "joy": 1,
+        "love": 2,
+        "anger": 3,
+        "fear": 4,
+        "surprise": 5,
+    }
     for i, k in enumerate(label2ids.keys()):
         label2ids[k] = outputs[0][i]
     label2ids = {k: float(v) for k, v in sorted(label2ids.items(), key=lambda item: item[1], reverse=True)}
@@ -35,7 +45,7 @@ def predict(query: str) -> dict:
 demo = gr.Interface(
     theme=gr.themes.Soft(),
     title="RHM Emotion Classifier 😊",
-    description="Beyond Words: Capturing the Essence of Emotion in Text<h3>On GPU it is much faster 🚀</h3>",
+    description="Beyond Words: Capturing the Essence of Emotion in Text<h3>On CPU with 8-bit quantization</h3>",
     fn=predict,
     inputs=gr.components.Textbox(label='Write your text here', lines=3),
     outputs=gr.components.Label(label='Predictions', num_top_classes=6),
@@ -51,6 +61,7 @@ demo.launch(share=True)
 
 
 
+
 #import gradio as gr
 # from transformers import pipeline, AutoTokenizer
 # from optimum.onnxruntime import ORTModelForSequenceClassification