Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,33 +1,41 @@
-import json, re, ast, streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-
-model_id = "mistralai/Mistral-7B-Instruct-v0.3"
-tok = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
-gen = pipeline("text-generation", model=model, tokenizer=tok,
-               max_new_tokens=256, do_sample=False, return_full_text=False)
-
-prompt = """Extract skills and knowledge from the text.
-Return JSON: {"SKILL":[...], "KNOWLEDGE":[...]}.
-Text: {text}
-JSON:"""
-
-def extract(text: str):
-    out = gen(prompt.format(text=text))
-    raw = out[0].get("generated_text") or out[0].get("text") or str(out[0])
-    m = re.search(r"\{[\s\S]*\}", raw)
-    data = {}
-    if m:
-        blob = m.group(0)
-        for parser in (json.loads, ast.literal_eval):
-            try:
-                data = parser(blob); break
-            except Exception: pass
-    if not isinstance(data, dict):
-        data = {}
-    return {"SKILL": data.get("SKILL", []), "KNOWLEDGE": data.get("KNOWLEDGE", [])}
-
-st.title("Skill/Knowledge Extractor")
-text = st.text_area("Paste text")
-if st.button("Extract") and text.strip():
+import json, re, ast, streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+
+model_id = "mistralai/Mistral-7B-Instruct-v0.3"
+tok = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
+gen = pipeline("text-generation", model=model, tokenizer=tok,
+               max_new_tokens=256, do_sample=False, return_full_text=False)
+
+prompt = """Extract skills and knowledge from the text.
+Return JSON: {"SKILL":[...], "KNOWLEDGE":[...]}.
+Text: {text}
+JSON:"""
+
+def extract(text: str):
+    out = gen(prompt.format(text=text))
+    raw = out[0].get("generated_text") or out[0].get("text") or str(out[0])
+    m = re.search(r"\{[\s\S]*\}", raw)
+    data = {}
+    if m:
+        blob = m.group(0)
+        for parser in (json.loads, ast.literal_eval):
+            try:
+                data = parser(blob); break
+            except Exception: pass
+    if not isinstance(data, dict):
+        return {
+            "SKILL": ["(Error: Invalid/Corrupted Model Output)"],
+            "KNOWLEDGE": [],
+            "DEBUG_RAW_OUTPUT": raw
+        }
+    
+    return {
+        "SKILL": data.get("SKILL", []),
+        "KNOWLEDGE": data.get("KNOWLEDGE", [])
+    }
+
+st.title("Skill/Knowledge Extractor")
+text = st.text_area("Paste text")
+if st.button("Extract") and text.strip():
     st.json(extract(text))