Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,6 @@
 import json, re, ast, streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch # Still needed for torch_dtype="auto" and device_map
+import torch 
 
 # SWITCHED MODEL: From Mistral-7B to the much smaller Gemma-2B-Instruct
 model_id = "google/gemma-2b-it"
@@ -12,7 +12,7 @@ tok = AutoTokenizer.from_pretrained(model_id)
 try:
     model = AutoModelForCausalLM.from_pretrained(
         model_id, 
-        torch_dtype=torch.bfloat16, # Use bfloat16 for better numerical stability if supported, otherwise auto
+        torch_dtype=torch.bfloat16, 
         device_map="auto"
     )
 except Exception: