Create app.py

app.py

@@ -0,0 +1,147 @@
+import os
+import torch
+import gradio as gr
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+
+# --------- CONFIG ---------
+BASE_MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+LORA_MODEL_ID = "nitya001/autotrain-fngb8-wqn4c"  # <-- change if your repo name differs
+
+MAX_NEW_TOKENS = 128
+TEMPERATURE = 0.7
+TOP_P = 0.9
+
+SYSTEM_PROMPT = (
+    "You are a helpful banking and loan support assistant. "
+    "You answer short, clear, and factual responses about UTRs, EMIs, "
+    "loan summaries, and payment issues based ONLY on the given question. "
+    "If you don't know something (like actual live data), say that you "
+    "cannot access real-time systems and answer generically."
+)
+
+# --------- LOAD MODEL ---------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(f"Loading base model: {BASE_MODEL_ID} on {device}...")
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+
+# TinyLlama uses eos_token as pad_token sometimes; ensure it's set
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_ID,
+    torch_dtype=torch.float32,
+    device_map=None,
+)
+
+print(f"Loading LoRA adapter: {LORA_MODEL_ID}...")
+model = PeftModel.from_pretrained(
+    base_model,
+    LORA_MODEL_ID,
+)
+model.to(device)
+model.eval()
+
+# --------- CHAT LOGIC ---------
+def format_chat_history(history, user_message):
+    """
+    Convert chat history + new user message into a single prompt string.
+    For now, we keep it simple: a system prompt + last few turns.
+    """
+    parts = [f"System: {SYSTEM_PROMPT}"]
+    for old_user, old_bot in history:
+        parts.append(f"User: {old_user}")
+        parts.append(f"Assistant: {old_bot}")
+    parts.append(f"User: {user_message}")
+    parts.append("Assistant:")
+    return "\n".join(parts)
+
+
+def generate_reply(user_message, history):
+    if not user_message.strip():
+        return history
+
+    # Build prompt from history
+    prompt = format_chat_history(history, user_message)
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=512,
+    ).to(device)
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=MAX_NEW_TOKENS,
+            do_sample=True,
+            temperature=TEMPERATURE,
+            top_p=TOP_P,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+
+    # Decode only the newly generated part
+    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # Naive way: take everything after the last "Assistant:" marker
+    if "Assistant:" in full_text:
+        bot_reply = full_text.split("Assistant:")[-1].strip()
+    else:
+        bot_reply = full_text.strip()
+
+    history.append((user_message, bot_reply))
+    return history
+
+
+# --------- GRADIO UI ---------
+with gr.Blocks(title="TinyLoan Assistant") as demo:
+    gr.Markdown(
+        """
+        # 💬 TinyLoan Assistant (TinyLlama + LoRA)
+        Ask about UTRs, EMIs, loan summaries, payment issues, etc.
+
+        > **Note:** This demo does not access real bank systems.  
+        > It answers based on patterns learned from example data.
+        """
+    )
+
+    chatbot = gr.Chatbot(
+        label="Chat",
+        height=400,
+        type="pairs",
+    )
+
+    with gr.Row():
+        user_input = gr.Textbox(
+            show_label=False,
+            placeholder="Type your question, e.g. 'What is my latest UTR?'",
+            scale=4,
+        )
+        send_btn = gr.Button("Send", scale=1)
+
+    clear_btn = gr.Button("Clear chat")
+
+    def respond(message, chat_history):
+        if chat_history is None:
+            chat_history = []
+        return generate_reply(message, chat_history)
+
+    send_btn.click(
+        respond,
+        inputs=[user_input, chatbot],
+        outputs=[chatbot],
+    )
+    user_input.submit(
+        respond,
+        inputs=[user_input, chatbot],
+        outputs=[chatbot],
+    )
+
+    clear_btn.click(lambda: [], outputs=[chatbot])
+
+demo.launch()