llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixes for runpod test

Browse Source
This commit is contained in:
geens
2024-11-16 21:51:18 +01:00
parent 9dfd5ddc39
commit c44d0ffbdf
8 changed files with 58 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
sia/local_llm_engine.py
View File

@@ -72,7 +72,7 @@ class LocalLlmEngine(LlmEngine):
text_inputs=prompt,
do_sample=True,
temperature=self._temperature,
max_new_tokens=self._token_limit,
max_new_tokens=self.token_limit(),
streamer=streamer
)
thread = Thread(target=self._pipeline, kwargs=pipeline_kwargs)
@@ -106,4 +106,7 @@ class LocalLlmEngine(LlmEngine):
Returns:
int: Maximum number of tokens the model can process
"""
return self._pipeline.model.config.max_position_embeddings
if self._token_limit is not None:
return self._token_limit
else:
return self._pipeline.model.config.max_position_embeddings

17
sia/repeat_entry.py
View File

@@ -99,19 +99,22 @@ class RepeatEntry(Entry):
element = ET.Element("repeat", {
"id": self.id,
})
if self._timeout:
element.set("timeout", str(self._timeout))
element.text = self.script
if self._exit_code is not None:
element.set("exit_code", str(self._exit_code))
if self._timed_out:
element.set("timed_out", "true")
else:
if len(self._stdout) > self._limit:
elif self._exit_code is not None:
element.set("exit_code", str(self._exit_code))
if self._limit:
element.set("limit", str(self._limit))
if len(self._stdout) > (self._limit or self.default_limit):
element.set("stdout_truncated", "true")
stdout_elem = ET.SubElement(element, "stdout")
stdout_elem.text = self._stdout[:self._limit]
if len(self._stderr) > self._limit:
stdout_elem.text = self._stdout[:(self._limit or self.default_limit)]
if len(self._stderr) > (self._limit or self.default_limit):
element.set("stderr_truncated", "true")
stderr_elem = ET.SubElement(element, "stderr")
stderr_elem.text = self._stderr[:self._limit]
stderr_elem.text = self._stderr[:(self._limit or self.default_limit)]
return element

7
sia/response_parser.py
View File

@@ -62,8 +62,13 @@ class ResponseParser:
entry_id = str(uuid.uuid4())
timestamp = datetime.now()
parser = ET.XMLPullParser(events=("start", "end"))
parser.feed(xml)
root = None
try:
root = ET.fromstring(xml.strip())
for event, root in parser.read_events():
if event == "start":
break
except ET.ParseError as e:
return ParseErrorEntry(entry_id, timestamp, xml, f"Invalid XML: {str(e)}")

14
sia/single_entry.py
View File

@@ -43,7 +43,7 @@ class SingleEntry(Entry):
super().__init__(id, timestamp)
self._script = script
self._timeout = timeout or self.default_timeout
self._limit = limit or self.default_limit
self._limit = limit
self._stdout = ""
self._stderr = ""
self._exit_code: Optional[int] = None
@@ -102,17 +102,21 @@ class SingleEntry(Entry):
element = ET.Element("single", {
"id": self.id,
})
if self._timeout:
element.set("timeout", str(self._timeout))
element.text = self.script
if self._timed_out:
element.set("timed_out", "true")
elif self._executed:
element.set("exit_code", str(self._exit_code))
if len(self._stdout) > self._limit:
if self._limit:
element.set("limit", str(self._limit))
if len(self._stdout) > (self._limit or self.default_limit):
element.set("stdout_truncated", "true")
stdout_elem = ET.SubElement(element, "stdout")
stdout_elem.text = self._stdout[:self._limit]
if len(self._stderr) > self._limit:
stdout_elem.text = self._stdout[:(self._limit or self.default_limit)]
if len(self._stderr) > (self._limit or self.default_limit):
element.set("stderr_truncated", "true")
stderr_elem = ET.SubElement(element, "stderr")
stderr_elem.text = self._stderr[:self._limit]
stderr_elem.text = self._stderr[:(self._limit or self.default_limit)]
return element