Spaces:
Sleeping
Sleeping
modified cpu calling
Browse files
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)
|