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

Added text streaming

Browse Source
This commit is contained in:
Niels Geens
2024-10-22 15:29:45 +02:00
parent 20429616f7
commit 4629e751dc
2 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
Dockerfile
View File

@@ -7,7 +7,7 @@ RUN pip3 install -r requirements.txt
FROM requirements AS test FROM requirements AS test
COPY ./test/ /root/test/ COPY ./test/ /root/test/
RUN mkdir -p /root/model RUN mkdir -p /root/model
CMD ["python3", "-m", "unittest", "discover", "-p", "*test.py", "-v"] CMD ["python3", "-m", "unittest", "discover", "-v", "-p", "*test.py", "-v"]
FROM requirements FROM requirements
CMD ["python3", "-m", "sia"] CMD ["python3", "-m", "sia"]

12
sia/llm_engine.py
View File

@@ -1,4 +1,4 @@
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, TextStreamer
import torch import torch
from . import util from . import util
@@ -34,12 +34,18 @@ class LlmEngine:
self.tokenizer.pad_token_id = self.tokenizer.eos_token_id self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
if model.config.pad_token_id is None: if model.config.pad_token_id is None:
model.config.pad_token_id = model.config.eos_token_id model.config.pad_token_id = model.config.eos_token_id
streamer = TextStreamer(
self.tokenizer,
skip_prompt=True
)
self.pipeline = pipeline( self.pipeline = pipeline(
"text-generation", "text-generation",
model=model, model=model,
tokenizer=self.tokenizer, tokenizer=self.tokenizer,
torch_dtype=torch.bfloat16, torch_dtype=torch.bfloat16,
device_map="auto", device_map="auto",
streamer=streamer,
return_full_text=False,
) )
def infer(self, system_prompt: str, main_context: str, action_schema: str) -> InferenceResult: def infer(self, system_prompt: str, main_context: str, action_schema: str) -> InferenceResult:
@@ -64,8 +70,8 @@ class LlmEngine:
) )
outputs = self.pipeline(prompt, max_new_tokens=120, do_sample=True) outputs = self.pipeline(prompt, max_new_tokens=120, do_sample=True)
generated_text = outputs[0]["generated_text"] generated_text = outputs[0]["generated_text"]
response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>",1)[1].strip() #response = generated_text.split("<|start_header_id|>assistant<|end_header_id|>",1)[1].strip()
result = util.split_response(response, valid_elements) result = util.split_response(generated_text, valid_elements)
return result return result
def finetune(self, dataset_paths: list, output_dir: str): def finetune(self, dataset_paths: list, output_dir: str):