Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 import torch
-from peft import PeftModel, PeftConfig
 from transformers import (
     AutoTokenizer,
     AutoModelForCausalLM,
@@ -10,34 +9,32 @@ from transformers import (
 # ---------------------------------------------------------------------
 # 1. Model Configuration
 # ---------------------------------------------------------------------
-ADAPTER_ID = "FrameRateTech/DamageScan-llama-8b-instruct-merged"
-
-# Load adapter config to find base model name
-peft_config = PeftConfig.from_pretrained(ADAPTER_ID)
-BASE_MODEL_ID = peft_config.base_model_name_or_path
+MODEL_ID = "FrameRateTech/DamageScan-llama-8b-instruct-merged"
 
 # ---------------------------------------------------------------------
 # 2. Load Tokenizer
 # ---------------------------------------------------------------------
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, use_fast=False)
+# For many LLaMA-based models, you often need use_fast=False and sometimes trust_remote_code=True
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    use_fast=False,
+    # trust_remote_code=True,   # Uncomment if needed for custom code
+)
+
+# Ensure we have a valid pad token
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
 # ---------------------------------------------------------------------
-# 3. Load Base Model + LoRA Weights
+# 3. Load the Model
 # ---------------------------------------------------------------------
-# If you need 8-bit to save VRAM, add load_in_8bit=True and device_map="auto"
-base_model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL_ID,
-    torch_dtype=torch.float16,
-    device_map="auto",
-)
-
-model = PeftModel.from_pretrained(
-    base_model,
-    ADAPTER_ID,
+# If you want to load 8-bit weights for VRAM savings, set load_in_8bit=True
+# and device_map="auto". Otherwise, below loads in FP16.
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
     torch_dtype=torch.float16,
     device_map="auto",
+    # trust_remote_code=True,   # Uncomment if needed for custom code
 )
 
 model.eval()
@@ -111,11 +108,10 @@ def predict(messages, temperature, top_p, max_new_tokens):
 # 7. Build the Gradio Interface
 # ---------------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("<h1 align='center'>FrameRateTech/DamageScan-llama-8b-instruct-merged Chatbot</h1>")
+    gr.Markdown("<h1 align='center'>DamageScan 8B Instruct Chatbot</h1>")
 
     with gr.Row():
         with gr.Column():
-            # type="messages" => each message is a dict with 'role' and 'content'
             chatbot = gr.Chatbot(label="Chat History", type="messages")
         with gr.Column():
             gr.Markdown("### Generation Settings")
@@ -129,17 +125,14 @@ with gr.Blocks() as demo:
                 minimum=64, maximum=2048, value=256, step=64, label="Max New Tokens"
             )
 
-    # Chat state is stored in 'chatbot' since type="messages"
     user_input = gr.Textbox(lines=1, label="Your Message", placeholder="Type here...")
     send_btn = gr.Button("Send")
 
-    # Append user input to chat, generate model reply, then clear input
     def user_submit(message_history, user_text, temp, top_p, max_tokens):
         message_history.append({"role": "user", "content": user_text})
         updated_messages = predict(message_history, temp, top_p, max_tokens)
         return updated_messages, ""
 
-    # Send button or pressing Enter triggers user_submit
     send_btn.click(
         user_submit,
         inputs=[chatbot, user_input, temperature_slider, top_p_slider, max_tokens_slider],
@@ -151,5 +144,4 @@ with gr.Blocks() as demo:
         outputs=[chatbot, user_input],
     )
 
-    # Launch the Gradio app
     demo.queue().launch()