Update app.py

app.py

@@ -1,57 +1,68 @@
+# 🚀 DigiTwin - TinyLLaMA ChatML Inference App
 import streamlit as st
 import torch
-import time
 import os
+import time
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from huggingface_hub import login
 
-# Hardcoded Hugging Face Token
-HF_TOKEN = os.environ.get("HF_TOKEN")  # or directly "hf_xxxxxx"
+# --- Hugging Face Token (use Streamlit secrets for secure deployment) ---
+HF_TOKEN = st.secrets["HF_TOKEN"]
 
-# App config
-st.set_page_config(
-    page_title="GM Fine-tune Assistant 🚀",
-    page_icon="🚀",
-    layout="centered"
-)
+# --- Streamlit Page Configuration ---
+st.set_page_config(page_title="DigiTwin - TinyLLaMA", page_icon="🦙", layout="centered")
 
-st.title("🚀 GM Fine-tune Assistant 🚀")
+# --- App Logo and Title ---
+# st.image("assets/valonylabz_logo.png", width=160)  # Optional: Add your logo
+st.title("🚀 DigiTwin - TinyLLaMA ChatML 🚀")
+
+# --- Model Path ---
+MODEL_ID = "amiguel/TinyLLaMA-110M-general-knowledge"
+
+# --- System Prompt (ChatML format) ---
+SYSTEM_PROMPT = (
+    "You are DigiTwin, the digital twin of Ataliba, an inspection engineer with over 17 years "
+    "of experience in mechanical integrity, reliability, piping, and asset management. "
+    "Be precise, practical, and technical. Provide advice aligned with industry best practices."
+)
 
-# Avatars
+# --- Avatars ---
 USER_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/9904d9a0d445ab0488cf7395cb863cce7621d897/USER_AVATAR.png"
 BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/991f4c6e4e1dc7a8e24876ca5aae5228bcdb4dba/Ataliba_Avatar.jpg"
 
-# Login to Huggingface
-login(token=HF_TOKEN)
-
-# Load Model
+# --- Load Model & Tokenizer ---
 @st.cache_resource
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("amiguel/GM_finetune", token=HF_TOKEN)
+def load_tinyllama_model():
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True,
+        use_fast=False,  # SentencePiece tokenizer
+        token=HF_TOKEN
+    )
     model = AutoModelForCausalLM.from_pretrained(
-        "amiguel/GM_finetune",
+        MODEL_ID,
         device_map="auto",
-        torch_dtype=torch.bfloat16,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        trust_remote_code=True,
         token=HF_TOKEN
     )
     return model, tokenizer
 
-model, tokenizer = load_model()
+model, tokenizer = load_tinyllama_model()
 
-# Session state
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+# --- Build ChatML Prompt ---
+def build_chatml_prompt(messages):
+    prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+    for msg in messages:
+        role = msg["role"]
+        prompt += f"<|im_start|>{role}\n{msg['content']}<|im_end|>\n"
+    prompt += "<|im_start|>assistant\n"
+    return prompt
 
-# Streamer
-def generate_response(prompt, model, tokenizer):
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True
-    )
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+# --- Generate Response with Streaming ---
+def generate_response(prompt_text):
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
 
     generation_kwargs = {
         "input_ids": inputs["input_ids"],
@@ -61,55 +72,53 @@ def generate_response(prompt, model, tokenizer):
         "top_p": 0.9,
         "repetition_penalty": 1.1,
         "do_sample": True,
-        "streamer": streamer
+        "streamer": streamer,
     }
 
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     return streamer
 
-# Display chat history
-for message in st.session_state.messages:
-    avatar = USER_AVATAR if message["role"] == "user" else BOT_AVATAR
-    with st.chat_message(message["role"], avatar=avatar):
-        st.markdown(message["content"])
+# --- Session State ---
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-# Chat input
-if prompt := st.chat_input("Ask me anything about General Knowledge..."):
+# --- Display Chat History ---
+for msg in st.session_state.messages:
+    avatar = USER_AVATAR if msg["role"] == "user" else BOT_AVATAR
+    with st.chat_message(msg["role"], avatar=avatar):
+        st.markdown(msg["content"])
 
-    # Display user message
+# --- Handle Chat Input ---
+if prompt := st.chat_input("Ask DigiTwin about inspection, piping, or reliability..."):
+    # Show user message
     with st.chat_message("user", avatar=USER_AVATAR):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Bot generating response
-    if model and tokenizer:
-        try:
-            with st.chat_message("assistant", avatar=BOT_AVATAR):
-                start_time = time.time()
-                streamer = generate_response(prompt, model, tokenizer)
-
-                response_container = st.empty()
-                full_response = ""
-
-                for chunk in streamer:
-                    full_response += chunk
-                    response_container.markdown(full_response + "▌", unsafe_allow_html=True)
-
-                end_time = time.time()
-                input_tokens = len(tokenizer(prompt)["input_ids"])
-                output_tokens = len(tokenizer(full_response)["input_ids"])
-                speed = output_tokens / (end_time - start_time)
-
-                st.caption(
-                    f"🔑 Input Tokens: {input_tokens} | Output Tokens: {output_tokens} | "
-                    f"🕒 Speed: {speed:.1f} tokens/sec"
-                )
-
-                response_container.markdown(full_response)
-                st.session_state.messages.append({"role": "assistant", "content": full_response})
-
-        except Exception as e:
-            st.error(f"⚡ Generation error: {str(e)}")
-    else:
-        st.error("🤖 Model not loaded!")
+    # Build ChatML-formatted prompt
+    prompt_text = build_chatml_prompt(st.session_state.messages)
+
+    with st.chat_message("assistant", avatar=BOT_AVATAR):
+        start = time.time()
+        streamer = generate_response(prompt_text)
+
+        response_area = st.empty()
+        full_response = ""
+
+        for chunk in streamer:
+            full_response += chunk.replace("<|im_end|>", "").strip() + " "
+            response_area.markdown(full_response + "▌", unsafe_allow_html=True)
+
+        end = time.time()
+        input_tokens = len(tokenizer(prompt_text)["input_ids"])
+        output_tokens = len(tokenizer(full_response)["input_ids"])
+        speed = output_tokens / (end - start)
+
+        st.caption(
+            f"🔑 Input Tokens: {input_tokens} | Output Tokens: {output_tokens} | "
+            f"🕒 Speed: {speed:.1f} tokens/sec"
+        )
+
+        response_area.markdown(full_response.strip())
+        st.session_state.messages.append({"role": "assistant", "content": full_response.strip()})