Rename api.py to fast_api.py

api.py → fast_api.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI, HTTPException
 from transformers import pipeline
 import uvicorn
+import streamlit as st
 
 # Load trained model
 model_name = "DINGOLANI/distilbert-ner-v2"
@@ -10,7 +11,7 @@ try:
 except Exception as e:
     raise RuntimeError(f"Failed to load model: {e}")
 
-# Label Mapping
+# Corrected label mapping based on expected training labels
 label_map = {
     "LABEL_1": "B-BRAND",
     "LABEL_2": "I-BRAND",
@@ -45,6 +46,8 @@ def predict(query: str):
         for label in result:
             label["score"] = float(label["score"])
 
+        print("RAW MODEL OUTPUT:", result)  
+
         structured_output = {}
         prev_label = None
         prev_word = None
@@ -60,7 +63,9 @@ def predict(query: str):
                     if prev_label == entity and prev_word:
                         structured_output[entity][-1] += word[2:]
                     else:
-                        structured_output.setdefault(entity, []).append(word)
+                        structured_output.setdefault(entity, []).append(word[2:])
+                else:
+                    structured_output.setdefault(entity, []).append(word)
 
                 prev_label = entity
                 prev_word = word
@@ -74,5 +79,28 @@ def predict(query: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error processing request: {e}")
 
+# 🚀 Streamlit Frontend
+def main():
+    st.set_page_config(page_title="Luxury Fashion NER", layout="wide")
+
+    st.title("👜 Luxury Fashion Entity Extractor")
+    st.write("Enter a text query and extract structured entities like **Brand, Category, Gender, and Price.**")
+
+    query = st.text_input("Enter Query:", "Gucci handbags for women under $5000")
+
+    if st.button("Analyze"):
+        response = predict(query)
+        
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            st.subheader("🔍 Structured Output")
+            for key, value in response["structured_output"].items():
+                st.write(f"**{key}:** {', '.join(value)}")
+        
+        with col2:
+            st.subheader("🛠 Raw Model Output")
+            st.json(response["raw_output"])
+
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    main()