Updated to load the model in 8-bit precision to reduce memory usage

app.py

@@ -2,7 +2,13 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 model_name = "ai4bharat/Airavata"
-model = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Load the model in 8-bit precision to reduce memory usage
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    device_map="auto", 
+    load_in_8bit=True
+)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 def generate_text(prompt, max_length):