Update app.py

app.py

@@ -1,62 +1,75 @@
-import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-# Pick one of the models from your collection:
-MODEL_NAME = "NextGLab/oransight-20-gemma-2b"  # <-- edit this to whichever in your collection
-
-@torch.inference_mode()
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        device_map="auto",
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
-    )
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=model.device,
-    )
-    return pipe
-
-pipe = load_model()
-
-# Simple chatbot fn
-def chat(message, history, max_new_tokens=256, temperature=0.7):
-    prompt = message
-    outputs = pipe(
-        prompt,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        do_sample=True,
-        pad_token_id=pipe.tokenizer.eos_token_id,
-    )
-    text = outputs[0]["generated_text"]
-    # Hugging Face pipeline often echos input text → remove
-    return text[len(prompt):].strip()
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Hugging Face repo ID (from the model page)
+MODEL_NAME = "NextGLab/ORANSight_Gemma_2_2B_Instruct"
+
+# Load tokenizer & model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
+
+# --- Helper function ---
+def chat(message, history, max_new_tokens=128, temperature=0.7):
+    """
+    message: user input
+    history: running chat history (list of [user, assistant])
+    """
+
+    # Convert Gradio-style history into chat template
+    messages = []
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": message})
+
+    # Prepare input using Gemma chat template
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt",
+    ).to(model.device)
 
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    # Decode only new tokens (avoid echoing input)
+    response = tokenizer.decode(
+        outputs[0][inputs["input_ids"].shape[-1]:],
+        skip_special_tokens=True
+    ).strip()
+
+    history.append((message, response))
+    return history, history, ""
+
+# --- Gradio App ---
 with gr.Blocks() as demo:
-    gr.Markdown(f"# 🤖 Oransight‑20‑Gemma Demo\nModel: **{MODEL_NAME}**")
-    
+    gr.Markdown("# 🤖 ORANSight Gemma 2 2B Instruct")
+
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your message")
+    msg = gr.Textbox(show_label=False, placeholder="Type a message...")
     send = gr.Button("Send")
-    clear = gr.Button("Clear")
-    
-    max_tokens = gr.Slider(50, 600, value=256, step=10, label="Max new tokens")
+    clear = gr.Button("Clear Chat")
+
+    max_tokens = gr.Slider(50, 512, value=128, step=10, label="Max new tokens")
     temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature")
-    
+
     state = gr.State([])
-    
-    def respond(message, history, max_tokens, temperature):
-        response = chat(message, history, max_tokens, temperature)
-        history.append((message, response))
-        return history, history, ""
-    
-    send.click(respond, [msg, state, max_tokens, temperature], [chatbot, state, msg])
-    msg.submit(respond, [msg, state, max_tokens, temperature], [chatbot, state, msg])
+
+    msg.submit(chat, [msg, state, max_tokens, temperature], [chatbot, state, msg])
+    send.click(chat, [msg, state, max_tokens, temperature], [chatbot, state, msg])
     clear.click(lambda: ([], []), None, [chatbot, state])
 
 if __name__ == "__main__":