aryan195a commited on
Commit
ca425df
·
1 Parent(s): 58c3420

modified cpu calling

Browse files
Files changed (1) hide show
  1. graph.py +2 -2
graph.py CHANGED
@@ -57,7 +57,7 @@ def generate(state: Dict[str, Any]) -> Dict[str, Any]:
57
  )
58
 
59
  try:
60
- response_raw = llm(
61
  prompt,
62
  temperature=temperature,
63
  do_sample=False,
@@ -105,7 +105,7 @@ def build_graph(retriever) -> Any:
105
  max_length=512,
106
  device=-1 # CPU only
107
  )
108
- llm = HuggingFacePipeline(pipeline=hf_pipe)
109
 
110
  graph = StateGraph(dict)
111
  graph.add_node("retrieve", retrieve)
 
57
  )
58
 
59
  try:
60
+ response_raw = llm.invoke(
61
  prompt,
62
  temperature=temperature,
63
  do_sample=False,
 
105
  max_length=512,
106
  device=-1 # CPU only
107
  )
108
+ llm = HuggingFacePipeline(pipeline=hf_pipe, model_kwargs={"device": -1})
109
 
110
  graph = StateGraph(dict)
111
  graph.add_node("retrieve", retrieve)