Deepakraj2006 commited on
Commit
5aacf46
·
verified ·
1 Parent(s): 1aeb2a1

Update worker.py

Browse files
Files changed (1) hide show
  1. worker.py +2 -1
worker.py CHANGED
@@ -30,7 +30,8 @@ def init_llm():
30
  if not hf_token:
31
  raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set in environment variables.")
32
 
33
- model_id = "mistralai/Mistral-7B-Instruct-v0.1"
 
34
  hf_pipeline = pipeline("text-generation", model=model_id, device=DEVICE)
35
  llm_pipeline = HuggingFacePipeline(pipeline=hf_pipeline)
36
 
 
30
  if not hf_token:
31
  raise ValueError("HUGGINGFACEHUB_API_TOKEN is not set in environment variables.")
32
 
33
+ model_id = model_id = "tiiuae/falcon-rw-1b" # Falcon-1B model
34
+
35
  hf_pipeline = pipeline("text-generation", model=model_id, device=DEVICE)
36
  llm_pipeline = HuggingFacePipeline(pipeline=hf_pipeline)
37