Remove context, change questions to prompt, and add upload select box

app.py

@@ -14,8 +14,9 @@ LICENSE = """
 # This is a dummy generation function
 import torch
 import gradio as gr
+import pandas as pd
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
+from peft import PeftMode
 
 # === [1] Model and Tokenizer Loading ===
 base_model_id = "NousResearch/Nous-Hermes-2-Mistral-7B-DPO"
@@ -32,22 +33,36 @@ model.eval()
 # 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.
+# 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:"""
+
+def create_prompt(table_file, prompt_text):
+    df = pd.read_csv(table_file.name)
+    markdown_table = df.to_markdown(index=False)
+    return f"""You are a financial assistant. Given the table and prompt, answer accordingly.
+
 Table:
-{table}
-Context:
-{context}
-Question:
-{question}
+{markdown_table}
+
+Prompt:
+{prompt_text}
+
 Answer:"""
 
-# === [3] Inference Function ===
-def answer_question(table, context, question):
-    prompt = create_prompt(table, context, question)
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)  # <-- use .to(device)
-    
+# Inference function
+def generate_answer(table_file, prompt_text):
+    full_prompt = create_prompt(table_file, prompt_text)
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -56,24 +71,21 @@ def answer_question(table, context, question):
             eos_token_id=tokenizer.eos_token_id
         )
 
-    # --- Decode and clean up ---
     output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     answer = output_text.split("Answer:")[-1].strip()
     return answer
 
-# === [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.")
+# Gradio UI layout
+with gr.Blocks(title="TAT-LLM Table & Prompt QA") as demo:
+    gr.Markdown("## TAT-LLM: Table & Prompt QA\nUpload your table (CSV format) and enter your prompt below.")
 
-    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?")
+    table_input = gr.File(label="Upload Table (.csv only)", file_types=[".csv"])
+    prompt_input = gr.Textbox(label="Prompt", lines=4, placeholder="Tuliskan pertanyaanmu di sini, termasuk konteks jika perlu.")
     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)
+    submit_btn.click(fn=generate_answer, inputs=[table_input, prompt_input], outputs=output_box)
 
-# === [5] Launch ===
+# Launch
 if __name__ == "__main__":
     demo.queue().launch()