Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,24 +1,27 @@
 import json, re, ast, streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
-import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch # Still needed for torch_dtype="auto" and device_map
 
-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
-)
+# SWITCHED MODEL: From Mistral-7B to the much smaller Gemma-2B-Instruct
+model_id = "google/gemma-2b-it"
 
 tok = AutoTokenizer.from_pretrained(model_id)
 
-# Load the model with the 4-bit quantization configuration
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    quantization_config=bnb_config,
-    device_map="auto"
-)
+# Simplified Model Loading: Removed BitsAndBytesConfig
+# This smaller model might load cleanly without 4-bit quantization, resolving the dependency issues.
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, 
+        torch_dtype=torch.bfloat16, # Use bfloat16 for better numerical stability if supported, otherwise auto
+        device_map="auto"
+    )
+except Exception:
+    # Fallback to float16 if bfloat16 causes issues
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, 
+        torch_dtype=torch.float16, 
+        device_map="auto"
+    )
 
 gen = pipeline("text-generation", model=model, tokenizer=tok,
                max_new_tokens=256, do_sample=False, return_full_text=False)
@@ -31,7 +34,6 @@ 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 = {}
     
@@ -40,7 +42,6 @@ def extract(text: str):
         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):
@@ -49,7 +50,6 @@ 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)"],
@@ -57,7 +57,6 @@ 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", [])
@@ -67,5 +66,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))