Update app.py

app.py

@@ -1,22 +1,11 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from huggingface_hub import login
-from threading import Thread
 import PyPDF2
 import pandas as pd
 import torch
-import time
 import os
 
-# Check if 'peft' is installed
-try:
-    from peft import PeftModel, PeftConfig
-except ImportError:
-    raise ImportError(
-        "The 'peft' library is required but not installed. "
-        "Please install it using: `pip install peft`"
-    )
-
 # Set page configuration
 st.set_page_config(
     page_title="WizNerd Insp",
@@ -25,14 +14,17 @@ st.set_page_config(
 )
 
 # Load Hugging Face token from environment variable
-HF_TOKEN = os.getenv("HF_TOKEN")  # Set this in your environment, e.g., via export HF_TOKEN="your_token"
+HF_TOKEN = os.getenv("HF_TOKEN")  # Set this in your Space's secrets
+
+# Model name
+MODEL_NAME = "amiguel/instruct_BERT-base-uncased_model"
 
-# Model names
-BASE_MODEL_NAME = "google-bert/bert-base-uncased"
-MODEL_OPTIONS = {
-    "Full Fine-Tuned": "amiguel/instruct_BERT-base-uncased_model",
-    "LoRA Adapter": "amiguel/SmolLM2-360M-concise-reasoning-lora",
-    "QLoRA Adapter": "amiguel/SmolLM2-360M-concise-reasoning-qlora"  # Hypothetical, adjust if needed
+# Label mapping (same as in Colab)
+LABEL_TO_CLASS = {
+    0: "Campaign", 1: "Corrosion Monitoring", 2: "Flare Tip", 3: "Flare TIP",
+    4: "FU Items", 5: "Intelligent Pigging", 6: "Lifting", 7: "Non Structural Tank",
+    8: "Piping", 9: "Pressure Safety Device", 10: "Pressure Vessel (VIE)",
+    11: "Pressure Vessel (VII)", 12: "Structure", 13: "Flame Arrestor"
 }
 
 # Title with rocket emojis
@@ -44,10 +36,6 @@ BOT_AVATAR = "https://raw.githubusercontent.com/achilela/vila_fofoka_analysis/99
 
 # Sidebar configuration
 with st.sidebar:
-    st.header("Model Selection 🤖")
-    model_type = st.selectbox("Choose Model Type", list(MODEL_OPTIONS.keys()), index=0)
-    selected_model = MODEL_OPTIONS[model_type]
-    
     st.header("Upload Documents 📂")
     uploaded_file = st.file_uploader(
         "Choose a PDF or XLSX file",
@@ -78,7 +66,7 @@ def process_file(uploaded_file):
 
 # Model loading function
 @st.cache_resource
-def load_model(hf_token, model_type, selected_model):
+def load_model(hf_token):
     try:
         if not hf_token:
             st.error("🔐 Authentication required! Please set the HF_TOKEN environment variable.")
@@ -86,34 +74,17 @@ def load_model(hf_token, model_type, selected_model):
         
         login(token=hf_token)
         
-        # Load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME, token=hf_token)
+        # Load tokenizer and model for classification
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=hf_token)
+        model = AutoModelForSequenceClassification.from_pretrained(
+            MODEL_NAME,
+            num_labels=len(LABEL_TO_CLASS),  # Ensure correct number of labels
+            token=hf_token
+        )
         
         # Determine device
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        
-        # Load model based on type
-        if model_type == "Full Fine-Tuned":
-            # Load full fine-tuned model directly
-            model = AutoModelForCausalLM.from_pretrained(
-                selected_model,
-                torch_dtype=torch.bfloat16,
-                token=hf_token
-            ).to(device)
-        else:
-            # Load base model and apply PEFT adapter
-            base_model = AutoModelForCausalLM.from_pretrained(
-                BASE_MODEL_NAME,
-                torch_dtype=torch.bfloat16,
-                token=hf_token
-            ).to(device)
-            model = PeftModel.from_pretrained(
-                base_model,
-                selected_model,
-                torch_dtype=torch.bfloat16,
-                is_trainable=False,  # Inference mode
-                token=hf_token
-            ).to(device)
+        model.to(device)
         
         return model, tokenizer
         
@@ -121,32 +92,22 @@ def load_model(hf_token, model_type, selected_model):
         st.error(f"🤖 Model loading failed: {str(e)}")
         return None
 
-# Generation function with KV caching
-def generate_with_kv_cache(prompt, file_context, model, tokenizer, use_cache=True):
-    full_prompt = f"Analyze this context:\n{file_context}\n\nQuestion: {prompt}\nAnswer:"
+# Classification function
+def classify_instruction(prompt, file_context, model, tokenizer):
+    full_prompt = f"Context:\n{file_context}\n\nInstruction: {prompt}"
     
-    streamer = TextIteratorStreamer(
-        tokenizer, 
-        skip_prompt=True, 
-        skip_special_tokens=True
-    )
+    model.eval()
+    device = model.device
     
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True, max_length=128)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
     
-    generation_kwargs = {
-        "input_ids": inputs["input_ids"],
-        "attention_mask": inputs["attention_mask"],
-        "max_new_tokens": 1024,
-        "temperature": 0.7,
-        "top_p": 0.9,
-        "repetition_penalty": 1.1,
-        "do_sample": True,
-        "use_cache": use_cache,
-        "streamer": streamer
-    }
+    with torch.no_grad():
+        outputs = model(**inputs)
+        prediction = outputs.logits.argmax().item()
+        class_name = LABEL_TO_CLASS[prediction]
     
-    Thread(target=model.generate, kwargs=generation_kwargs).start()
-    return streamer
+    return class_name
 
 # Display chat messages
 for message in st.session_state.messages:
@@ -160,15 +121,14 @@ for message in st.session_state.messages:
 
 # Chat input handling
 if prompt := st.chat_input("Ask your inspection question..."):
-    # Load model if not already loaded or if model type changed
-    if "model" not in st.session_state or st.session_state.get("model_type") != model_type:
-        model_data = load_model(HF_TOKEN, model_type, selected_model)
+    # Load model if not already loaded
+    if "model" not in st.session_state:
+        model_data = load_model(HF_TOKEN)
         if model_data is None:
             st.error("Failed to load model. Please ensure HF_TOKEN is set correctly.")
             st.stop()
             
         st.session_state.model, st.session_state.tokenizer = model_data
-        st.session_state.model_type = model_type
     
     model = st.session_state.model
     tokenizer = st.session_state.tokenizer
@@ -178,47 +138,19 @@ if prompt := st.chat_input("Ask your inspection question..."):
         st.markdown(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
 
-    # Process Rank
+    # Process file context
     file_context = process_file(uploaded_file)
     
-    # Generate response with KV caching
+    # Classify the instruction
     if model and tokenizer:
         try:
             with st.chat_message("assistant", avatar=BOT_AVATAR):
-                start_time = time.time()
-                streamer = generate_with_kv_cache(prompt, file_context, model, tokenizer, use_cache=True)
-                
-                response_container = st.empty()
-                full_response = ""
-                
-                for chunk in streamer:
-                    cleaned_chunk = chunk.replace("<think>", "").replace("</think>", "").strip()
-                    full_response += cleaned_chunk + " "
-                    response_container.markdown(full_response + "▌", unsafe_allow_html=True)
-                
-                # Calculate performance metrics
-                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)
-                
-                # Calculate costs (hypothetical pricing model)
-                input_cost = (input_tokens / 1000000) * 5  # $5 per million input tokens
-                output_cost = (output_tokens / 1000000) * 15  # $15 per million output tokens
-                total_cost_usd = input_cost + output_cost
-                total_cost_aoa = total_cost_usd * 1160  # Convert to AOA (Angolan Kwanza)
-                
-                # Display metrics
-                st.caption(
-                    f"🔑 Input Tokens: {input_tokens} | Output Tokens: {output_tokens} | "
-                    f"🕒 Speed: {speed:.1f}t/s | 💰 Cost (USD): ${total_cost_usd:.4f} | "
-                    f"💵 Cost (AOA): {total_cost_aoa:.4f}"
-                )
-                
-                response_container.markdown(full_response)
-                st.session_state.messages.append({"role": "assistant", "content": full_response})
+                predicted_class = classify_instruction(prompt, file_context, model, tokenizer)
+                response = f"Predicted class: {predicted_class}"
+                st.markdown(response)
+                st.session_state.messages.append({"role": "assistant", "content": response})
                 
         except Exception as e:
-            st.error(f"⚡ Generation error: {str(e)}")
+            st.error(f"⚡ Classification error: {str(e)}")
     else:
         st.error("🤖 Model not loaded!")