classfinetune / app.py
amiguel's picture
Update app.py
33be14e verified
raw
history blame
5.37 kB
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import login
import PyPDF2
import pandas as pd
import torch
import os
# Set page configuration
st.set_page_config(
page_title="WizNerd Insp",
page_icon="πŸš€",
layout="centered"
)
# Load Hugging Face token from environment variable
HF_TOKEN = os.getenv("HF_TOKEN") # Set this in your Space's secrets
# Model name
MODEL_NAME = "amiguel/instruct_BERT-base-uncased_model"
# 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
st.title("πŸš€ WizNerd Insp πŸš€")
# Configure 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"
# Sidebar configuration
with st.sidebar:
st.header("Upload Documents πŸ“‚")
uploaded_file = st.file_uploader(
"Choose a PDF or XLSX file",
type=["pdf", "xlsx"],
label_visibility="collapsed"
)
# Initialize chat history
if "messages" not in st.session_state:
st.session_state.messages = []
# File processing function
@st.cache_data
def process_file(uploaded_file):
if uploaded_file is None:
return ""
try:
if uploaded_file.type == "application/pdf":
pdf_reader = PyPDF2.PdfReader(uploaded_file)
return "\n".join([page.extract_text() for page in pdf_reader.pages])
elif uploaded_file.type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
df = pd.read_excel(uploaded_file)
return df.to_markdown()
except Exception as e:
st.error(f"πŸ“„ Error processing file: {str(e)}")
return ""
# Model loading function
@st.cache_resource
def load_model(hf_token):
try:
if not hf_token:
st.error("πŸ” Authentication required! Please set the HF_TOKEN environment variable.")
return None
login(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"
model.to(device)
return model, tokenizer
except Exception as e:
st.error(f"πŸ€– Model loading failed: {str(e)}")
return None
# Classification function
def classify_instruction(prompt, file_context, model, tokenizer):
full_prompt = f"Context:\n{file_context}\n\nInstruction: {prompt}"
model.eval()
device = 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()}
with torch.no_grad():
outputs = model(**inputs)
prediction = outputs.logits.argmax().item()
class_name = LABEL_TO_CLASS[prediction]
return class_name
# Display chat messages
for message in st.session_state.messages:
try:
avatar = USER_AVATAR if message["role"] == "user" else BOT_AVATAR
with st.chat_message(message["role"], avatar=avatar):
st.markdown(message["content"])
except:
with st.chat_message(message["role"]):
st.markdown(message["content"])
# Chat input handling
if prompt := st.chat_input("Ask your inspection question..."):
# 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
model = st.session_state.model
tokenizer = st.session_state.tokenizer
# Add user message
with st.chat_message("user", avatar=USER_AVATAR):
st.markdown(prompt)
st.session_state.messages.append({"role": "user", "content": prompt})
# Process file context
file_context = process_file(uploaded_file)
# Classify the instruction
if model and tokenizer:
try:
with st.chat_message("assistant", avatar=BOT_AVATAR):
predicted_class = classify_instruction(prompt, file_context, model, tokenizer)
response = f"The Item Class is: {predicted_class}"
st.markdown(response)
st.session_state.messages.append({"role": "assistant", "content": response})
except Exception as e:
st.error(f"⚑ Classification error: {str(e)}")
else:
st.error("πŸ€– Model not loaded!")