Initial update for our case

app.py

@@ -1,6 +1,4 @@
 # chatbot_template.py
-
-import gradio as gr
 import spaces
 
 DESCRIPTION = """
@@ -16,50 +14,64 @@ LICENSE = """
 
 # This is a dummy generation function
 @spaces.GPU  # This allows it to run on GPU Spaces (remove if not needed)
-def generate_response(
-    message: str,
-    chat_history: list[dict],
-    system_prompt: str = "",
-    max_new_tokens: int = 512,
-    temperature: float = 0.7,
-    top_p: float = 0.9,
-    top_k: int = 40,
-    repetition_penalty: float = 1.1,
-):
-    # Replace this with actual model logic
-    yield f"(This is a dummy response to): {message}"
-
-    # Example for real model use:
-    # tokenizer = AutoTokenizer.from_pretrained("your-model-id")
-    # model = AutoModelForCausalLM.from_pretrained("your-model-id", device_map="auto")
-    # input_ids = tokenizer(message, return_tensors="pt").input_ids.to(model.device)
-    # output = model.generate(input_ids, ...)
-    # response = tokenizer.decode(output[0], skip_special_tokens=True)
-    # yield response
-
-chat_interface = gr.ChatInterface(
-    fn=generate_response,
-    additional_inputs=[
-        gr.Textbox(label="System Prompt", lines=2, value=""),
-        gr.Slider(label="Max new tokens", minimum=16, maximum=2048, value=512, step=16),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7, step=0.1),
-        gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, value=0.9, step=0.05),
-        gr.Slider(label="Top-k", minimum=1, maximum=1000, value=40, step=1),
-        gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, value=1.1, step=0.05),
-    ],
-    examples=[
-        ["Hello!"],
-        ["Can you summarize AI in one sentence?"],
-        ["What is the capital of France?"],
-    ],
-    cache_examples=False,
-    type="messages",
-)
-
-with gr.Blocks(fill_height=True) as demo:
-    gr.Markdown(DESCRIPTION)
-    chat_interface.render()
-    gr.Markdown(LICENSE)
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+
+# === [1] Model and Tokenizer Loading ===
+base_model_id = "meta-llama/Llama-2-7b-hf"  # Replace with your base model
+lora_path = "./tat-llm"  # Path to your fine-tuned LoRA folder
+
+# Load base model and LoRA adapter
+base_model = AutoModelForCausalLM.from_pretrained(base_model_id, torch_dtype=torch.float16)
+model = PeftModel.from_pretrained(base_model, lora_path)
+model.eval().cuda()
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(lora_path)
+
+# === [2] Prompt Formatting Function ===
+def create_prompt(table, context, question):
+    return f"""You are a financial assistant. Given the table and context, answer the question.
+
+Table:
+{table}
+
+Context:
+{context}
+
+Question:
+{question}
+
+Answer:"""
+
+# === [3] Inference Function ===
+def answer_question(table, context, question):
+    prompt = create_prompt(table, context, question)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=128,
+            do_sample=False,
+            eos_token_id=tokenizer.eos_token_id
+        )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+# === [4] Gradio UI Layout ===
+with gr.Blocks(title="TAT-LLM Table & Text QA") as demo:
+    gr.Markdown("## TAT-LLM: Table-and-Text Question Answering\nUpload a table (Markdown format), provide context, and ask your question.")
+
+    with gr.Row():
+        table_input = gr.Textbox(label="Table (Markdown)", lines=10, placeholder="| Quarter | Revenue |\n|--------|---------|\n| Q1 | 100 | ...")
+        context_input = gr.Textbox(label="Context", lines=10, placeholder="PT ABC mengalami peningkatan pendapatan dari Q1 ke Q4.")
+    question_input = gr.Textbox(label="Question", lines=2, placeholder="Berapa persentase kenaikan dari Q1 ke Q4?")
+    output_box = gr.Textbox(label="Answer", lines=5)
+
+    submit_btn = gr.Button("Generate Answer")
+    submit_btn.click(fn=answer_question, inputs=[table_input, context_input, question_input], outputs=output_box)
 
+# === [5] Launch ===
 if __name__ == "__main__":
     demo.queue().launch()