Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,9 +1,25 @@
 import json, re, ast, streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
+import torch
 
 model_id = "mistralai/Mistral-7B-Instruct-v0.3"
+
+# 4-bit Quantization Configuration to reduce memory usage (VRAM/RAM)
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16
+)
+
 tok = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
+
+# Load the model with the 4-bit quantization configuration
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config=bnb_config,
+    device_map="auto"
+)
+
 gen = pipeline("text-generation", model=model, tokenizer=tok,
                max_new_tokens=256, do_sample=False, return_full_text=False)
 
@@ -15,13 +31,16 @@ 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])
+    # Relaxed regex to find JSON object anywhere in the output
     m = re.search(r"(\{[\s\S]*\})", raw)
     data = {}
+    
     if m:
         blob = m.group(0).strip()
         for parser in (json.loads, ast.literal_eval):
             try:
                 parsed_data = parser(blob)
+                # Handle common case where model returns a list of dictionaries
                 if isinstance(parsed_data, list) and parsed_data:
                     data = parsed_data[0]
                 elif isinstance(parsed_data, dict):
@@ -30,6 +49,7 @@ def extract(text: str):
             except Exception:
                 continue
 
+    # Error handling check: If parsing failed completely, return a structured error dictionary
     if not isinstance(data, dict):
         return {
             "SKILL": ["(Error: Invalid/Corrupted Model Output)"],
@@ -37,6 +57,7 @@ def extract(text: str):
             "DEBUG_RAW_OUTPUT": raw
         }
 
+    # Successful return: Uses .get() which prevents KeyError even if keys are missing
     return {
         "SKILL": data.get("SKILL", []),
         "KNOWLEDGE": data.get("KNOWLEDGE", [])
@@ -46,5 +67,4 @@ st.title("Skill/Knowledge Extractor")
 text = st.text_area("Paste text")
 
 if st.button("Extract") and text.strip():
-    st.json(extract(text))
-    
+    st.json(extract(text))