halimbahae commited on
Commit
04a4930
·
verified ·
1 Parent(s): d4bd433

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -6
app.py CHANGED
@@ -20,12 +20,17 @@ if "scrape_results" not in st.session_state:
20
  # Load Zephyr model
21
  @st.cache_resource
22
  def load_model():
23
- return pipeline(
24
- "text-generation",
25
- model="HuggingFaceH4/zephyr-7b-beta",
26
- torch_dtype=torch.float16,
27
- device_map="auto",
28
- )
 
 
 
 
 
29
 
30
  # Initialize the model
31
  model = load_model()
 
20
  # Load Zephyr model
21
  @st.cache_resource
22
  def load_model():
23
+ try:
24
+ return pipeline(
25
+ "text-generation",
26
+ model="HuggingFaceH4/zephyr-7b-beta",
27
+ torch_dtype=torch.float16,
28
+ device_map="auto",
29
+ model_kwargs={"load_in_8bit": True} # Use 8-bit quantization to reduce memory usage
30
+ )
31
+ except Exception as e:
32
+ st.error(f"Error loading model: {str(e)}")
33
+ return None
34
 
35
  # Initialize the model
36
  model = load_model()