Update app.py

app.py

@@ -1,9 +1,31 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-def predict(text):
-    return f"Echo: {text}"
+model_name = config["model_name"]
 
-iface = gr.Interface(fn=predict, inputs="text", outputs="text")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, low_cpu_mem_usage=True, device_map="auto", torch_dtype="auto"
+)
+
+def predict(messages):
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    logger.info(f"Model generation process started at - {process_id}")
+    generated_ids = model.generate(**model_inputs, max_new_tokens=512)
+    generated_ids = [
+        output_ids[len(input_ids) :]
+        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+    logger.info(f"Model generation process completed [{process_id}]")
+
+    reply = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return reply
+
+iface = gr.Interface(fn=predict, inputs="messages", outputs="reply")
 
 # Launch with API access
-iface.launch(server_name="0.0.0.0", server_port=7860, share=True)
+iface.launch(server_name="0.0.0.0", server_port=7860, share=False)