Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import accelerate
 
 # Load the model and tokenizer
 model_name = "Telugu-LLM-Labs/Indic-gemma-2b-finetuned-sft-Navarasa-2.0"
-model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=False, device_map="auto")
+accelerator = accelerate.Accelerator()
+model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=False, device_map="auto", offload_folder="/tmp")
+model = accelerator.prepare(model)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 def generate_prompt(instruction, user_input):
@@ -26,8 +29,9 @@ def get_model_response(user_input, instruction="Identify and summarize the core
     """
     input_text = generate_prompt(instruction, user_input)
     inputs = tokenizer([input_text], return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=300, use_cache=True)
-    response = tokenizer.batch_decode(outputs)[0]
+    with accelerator.distribute_inputs_to_prepared(model.device_map, inputs):
+        outputs = model.generate(**inputs, max_new_tokens=300, use_cache=True)
+    response = tokenizer.batch_decode(accelerator.gather(outputs))[0]
     return response.split("### Response:")[-1].strip()
 
 # Gradio interface