Update app.py

app.py

@@ -1,16 +1,35 @@
 import spaces
 import gradio as gr
 
-from airllm import HuggingFaceModelLoader, AutoModelForCausalLM
+# from airllm import HuggingFaceModelLoader, AutoModelForCausalLM
+
+from airllm import AutoModel
+import mlx.core as mx
+
+model = AutoModel("meta-llama/Meta-Llama-3-8B-Instruct")
+# model = AutoModel.from_pretrained(model_loader)
+MAX_LENGTH = 128 
 
-model_loader = HuggingFaceModelLoader("meta-llama/Meta-Llama-3-8B-Instruct")
-model = AutoModelForCausalLM.from_pretrained(model_loader)
 
 @spaces.GPU
 def generate_text(input_text):
-    input_ids = model.tokenizer.encode(input_text, return_tensors="pt")
-    output = model.generate(input_ids, max_length=100)
-    return model.tokenizer.decode(output[0])
+    
+    input_tokens = model.tokenizer(input_text,
+                                   return_tensors="np", 
+                                   return_attention_mask=False, 
+                                   truncation=True, 
+                                   max_length=MAX_LENGTH, 
+                                   padding=False)
+
+               
+    output = model.generate(mx.array(input_tokens['input_ids']),
+                                       max_new_tokens=20,
+                                       use_cache=True,
+                                       return_dict_in_generate=True)
+    # input_ids = model.tokenizer.encode(input_text, return_tensors="np")
+    # output = model.generate(input_ids, max_length=100)
+    # return model.tokenizer.decode(output[0])
+    return output
 
 
 iface = gr.Interface(