Update app.py

app.py

@@ -4,6 +4,9 @@ from PIL import Image
 from deepseek_vl2.serve.inference import load_model, deepseek_generate, convert_conversation_to_prompts
 from deepseek_vl2.serve.app_modules.utils import configure_logger, strip_stop_words, pil_to_base64
 
+# Set Page Config (Must be first!)
+st.set_page_config(layout="wide")
+
 # Set up logging
 logger = configure_logger()
 
@@ -20,7 +23,7 @@ def fetch_model(model_name: str, dtype=torch.bfloat16):
         model_info = load_model(model_name, dtype=dtype)
         tokenizer, model, vl_chat_processor = model_info
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-        model = model.to(device).eval()  # Move to appropriate device and set eval mode
+        model = model.to(device).eval()  # Move to appropriate device
         DEPLOY_MODELS[model_name] = (tokenizer, model, vl_chat_processor)
         logger.info(f"Loaded {model_name} on {device}")
     return DEPLOY_MODELS[model_name]
@@ -52,22 +55,21 @@ def to_gradio_chatbot(conv):
             ret[-1][-1] = msg
     return ret
 
-# Predict function
-def predict(text, images, chatbot, history, model_name="deepseek-ai/deepseek-vl2-tiny"):
+# Predict function (simplified for OCR)
+def predict(text, images, model_name="deepseek-ai/deepseek-vl2-tiny"):
     logger.info("Starting predict function...")
     tokenizer, vl_gpt, vl_chat_processor = fetch_model(model_name)
     if not text:
         logger.warning("Empty text input detected.")
-        return chatbot, history, "Empty context."
+        return "Empty context."
 
     logger.info("Processing images...")
     pil_images = [Image.open(img).convert("RGB") for img in images] if images else []
     conversation = generate_prompt_with_history(
-        text, pil_images, history, vl_chat_processor, tokenizer
+        text, pil_images, [], vl_chat_processor, tokenizer
     )
     all_conv, _ = convert_conversation_to_prompts(conversation)
     stop_words = conversation.stop_str
-    gradio_chatbot_output = to_gradio_chatbot(conversation)
 
     full_response = ""
     logger.info("Generating response...")
@@ -85,53 +87,96 @@ def predict(text, images, chatbot, history, model_name="deepseek-ai/deepseek-vl2
                 repetition_penalty=1.1
             ):
                 full_response += x
-                response = strip_stop_words(full_response, stop_words)
-                conversation.update_last_message(response)
-                gradio_chatbot_output[-1][1] = response
-                logger.info(f"Yielding partial response: {response[:50]}...")
-                yield gradio_chatbot_output, conversation.messages, "Generating..."
-
+        response = strip_stop_words(full_response, stop_words)
         logger.info("Generation complete.")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()  # Clear cache only if CUDA is available
-        yield gradio_chatbot_output, conversation.messages, "Success"
+        return response
     except Exception as e:
         logger.error(f"Error in generation: {str(e)}")
-        yield gradio_chatbot_output, conversation.messages, f"Error: {str(e)}"
+        return f"Error: {str(e)}"
 
-# Streamlit OCR app interface
-def upload_and_process(image):
+# OCR extraction function
+def extract_text(image):
     if image is None:
-        return "Please upload an image.", []
+        return "Please upload an image."
     prompt = "Extract all text from this image exactly as it appears, ensuring the output is in English only. Preserve spaces, bullets, numbers, and all formatting. Do not translate, generate, or include text in any other language. Stop at the last character of the image text."
-    chatbot = []
-    history = []
-    logger.info("Starting upload_and_process...")
-    for chatbot_output, history_output, status in predict(prompt, [image], chatbot, history):
-        logger.info(f"Status: {status}")
-        if status == "Success":
-            return chatbot_output[-1][1], history_output
-        elif status.startswith("Error"):
-            return f"Error: {status}", []
-    return "Processing failed.", []
-
-# Streamlit UI
-st.markdown("<h1 style='text-align: center;'>🔍 Extract Job Info in One Click</h1>", unsafe_allow_html=True)
-
-uploaded_file = st.file_uploader("Upload an Image (PNG, JPG, JPEG)", type=["png", "jpg", "jpeg"])
-
-if uploaded_file:
-    extracted_text, _ = upload_and_process(uploaded_file)
-
-    col1, col2 = st.columns(2)
-    with col1:
-        st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
-    with col2:
-        st.markdown(
-            f"""
-            <div style="border: 1px solid #ccc; padding: 10px; width: 100%; white-space: pre-wrap; overflow: hidden; text-align: left;">
-                {extracted_text}
-            </div>
-            """,
-            unsafe_allow_html=True
-        )
+    logger.info("Starting text extraction...")
+    extracted_text = predict(prompt, [image])
+    return extracted_text
+
+# User Selection State
+if "user_selected" not in st.session_state:
+    st.session_state["user_selected"] = False
+if "selected_user" not in st.session_state:
+    st.session_state["selected_user"] = None
+
+# User Selection Page
+@st.cache_data
+def get_user_names():
+    return ["DS", "NW", "RB", "IG", "AR", "NU"]
+
+user_names = get_user_names()
+
+# Retrieve cached user from query params
+query_params = st.query_params
+cached_user = query_params.get("user", None)
+
+if cached_user and "user_selected" not in st.session_state:
+    st.session_state["user_selected"] = True
+    st.session_state["selected_user"] = cached_user
+
+if not st.session_state["user_selected"]:
+    st.title("👋 Let’s Get Started! Identify Yourself to Begin")
+    selected_user = st.selectbox("Choose your name:", user_names, index=user_names.index(cached_user) if cached_user in user_names else None)
+    continue_button = st.button("Continue", disabled=not selected_user)
+
+    if continue_button:
+        st.session_state["user_selected"] = True
+        st.session_state["selected_user"] = selected_user
+        st.query_params["user"] = selected_user
+        st.rerun()
+
+    if not selected_user:
+        st.warning("⚠ Please select a user to continue.")
+else:
+    st.write(f"✅ Welcome Back, {st.session_state['selected_user']}!")
+
+# Main UI (Only loads after user selection)
+if st.session_state["user_selected"]:
+    st.markdown("<h1 style='text-align: center;'>🔍 Extract Job Info in One Click</h1>", unsafe_allow_html=True)
+
+    uploaded_file = st.file_uploader("Upload an Image (PNG, JPG, JPEG)", type=["png", "jpg", "jpeg"])
+
+    if uploaded_file:
+        img = Image.open(uploaded_file)
+        extracted_text = extract_text(uploaded_file)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.image(uploaded_file, caption="Uploaded Image", use_container_width=True)
+        with col2:
+            st.markdown(
+                f"""
+                <div style="border: 1px solid #ccc; padding: 10px; width: 100%; white-space: pre-wrap; overflow: hidden; text-align: left;">
+                    {extracted_text}
+                </div>
+                """,
+                unsafe_allow_html=True
+            )
+
+        st.session_state["extracted_text"] = extracted_text
+
+    errors_text = st.text_area("Paste Any Errors Here", height=200)
+    rating = st.radio("Rate the OCR Extraction (5-1)", options=[5, 4, 3, 2, 1], index=None, horizontal=True)
+    ref_number = st.text_input("Enter Reference Number", max_chars=10)
+
+    if ref_number and not ref_number.isdigit():
+        st.warning("⚠ Reference Number must be a number.")
+        ref_number = ""
+
+    if not ref_number or rating is None:
+        st.warning("⚠ Please enter a Reference Number and select a Rating to proceed.")
+        st.button("Submit", disabled=True)
+    else:
+        if st.button("Submit"):
+            st.success("✅ Submitted successfully!")
+            # No Google Drive or Sheets upload; just a confirmation