Update app.py

app.py

@@ -1,59 +1,71 @@
 import streamlit as st
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-import torch
-
-model_name = "amiguel/fintune_naming_model"  # Replace with your model repo
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-def classify_review(text, model, tokenizer, device, max_length=512):
-    model.eval()
-    inputs = tokenizer.encode_plus(
-        text,
-        truncation=True,
-        padding='max_length',
-        max_length=max_length,
-        return_tensors="pt"
-    )
-    input_ids = inputs['input_ids'].to(device)
-    attention_mask = inputs['attention_mask'].to(device)
-
-    with torch.no_grad():
-        outputs = model(input_ids, attention_mask=attention_mask)
-        logits = outputs.logits
-    predicted_label = torch.argmax(logits, dim=-1).item()
-    return "Proper Naming otfcn" if predicted_label == 1 else "Wrong Naming notfcn"
-
-def main():
-    st.title("Notifications Naming Classifier")
-
-    input_option = st.radio("Select input option", ("Single Text Query", "Upload Table"))
-
-    if input_option == "Single Text Query":
-        text_query = st.text_input("Enter text query")
-        if st.button("Classify"):
-            if text_query:
-                predicted_label = classify_review(text_query, model, tokenizer, device)
-                st.write("Predicted Label:")
-                st.write(predicted_label)
-            else:
-                st.warning("Please enter a text query.")
-
-    elif input_option == "Upload Table":
-        uploaded_file = st.file_uploader("Choose a file", type=["csv", "xlsx"])
-        if uploaded_file is not None:
-            import pandas as pd
-            if uploaded_file.name.endswith(".csv"):
-                df = pd.read_csv(uploaded_file)
-            else:
-                df = pd.read_excel(uploaded_file)
-
-            text_column = st.selectbox("Select the text column", df.columns)
-            predicted_labels = [classify_review(text, model, tokenizer, device) for text in df[text_column]]
-            df["Predicted Label"] = predicted_labels
-            st.write(df) 
-
-if __name__ == "__main__":
-    main()
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from PyPDF2 import PdfReader
+import pandas as pd
+
+# Title and emojis
+st.title("🚀 WizNerd Insp 🚀")
+
+# Sidebar for file uploads
+st.sidebar.header("Upload Files")
+uploaded_xlsx = st.sidebar.file_uploader("Upload XLSX File", type=["xlsx"])
+uploaded_pdf = st.sidebar.file_uploader("Upload PDF File", type=["pdf"])
+
+# Load the HuggingFace model and tokenizer
+@st.cache_resource
+def load_model():
+    model_name = "amiguel/optimizedModelLinsting6.1"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+# Prompt style
+prompt_style = """
+Below is an instruction that describes a task, paired with an input that provides further context.
+Write a response that appropriately completes the request.
+Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.
+### Instruction:
+You are an experienced inspection methods engineer, a topside expert with advanced knowledge in scope definition, functional location determination, and inspection plan building.
+Please answer the following inspection scope question.
+### Instruction:
+{}
+### Output:
+<think>   {} </think> {}
+"""
+
+# Function to process user input and generate response
+def generate_response(input_text):
+    # Format the input using the prompt style
+    formatted_input = prompt_style.format(input_text, "", "")
+    
+    # Tokenize and generate response
+    inputs = tokenizer(formatted_input, return_tensors="pt", truncation=True, max_length=512)
+    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return response
+
+# Main chat interface
+st.header("Chat with WizNerd Insp")
+user_input = st.text_input("Ask a question:")
+if st.button("Submit"):
+    if user_input.strip() != "":
+        response = generate_response(user_input)
+        st.write("Response:")
+        st.write(response)
+
+# Process uploaded files
+if uploaded_xlsx:
+    st.write("Processing XLSX file...")
+    df = pd.read_excel(uploaded_xlsx)
+    st.write(df)
+
+if uploaded_pdf:
+    st.write("Processing PDF file...")
+    pdf_reader = PdfReader(uploaded_pdf)
+    text = ""
+    for page in pdf_reader.pages:
+        text += page.extract_text()
+    st.write(text)