github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-05-30 11:39:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

Nvidia trt model name for stop_stream() (#16997)

Browse Source 
just removing some legacy leftover.
This commit is contained in:
Mikhail Khludnev 2024-02-05 21:45:06 +03:00 committed by GitHub
parent 2ef69fe11b
commit 2145636f1d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
libs/partners/nvidia-trt/langchain_nvidia_trt/llms.py
View File

@ -377,14 +377,14 @@ class StreamingResponseGenerator(queue.Queue):
def __init__(
self,
client: grpcclient.InferenceServerClient,
llm: TritonTensorRTLLM,
request_id: str,
force_batch: bool,
stop_words: Sequence[str],
) -> None:
"""Instantiate the generator class."""
super().__init__()
self.client = client
self.llm = llm
self.request_id = request_id
self._batch = force_batch
self._stop_words = stop_words
@ -397,8 +397,8 @@ class StreamingResponseGenerator(queue.Queue):
"""Return the next retrieved token."""
val = self.get()
if val is None or val in self._stop_words:
self.client.stop_stream(
"tensorrt_llm", self.request_id, signal=not self._batch
self.llm.stop_stream(
self.llm.model_name, self.request_id, signal=not self._batch
)
raise StopIteration()
return val