Repository navigation
Expand file tree
/
Copy pathbase_llm.py
More file actions
93 lines (76 loc) · 2.93 KB
/
Copy pathbase_llm.py
File metadata and controls
93 lines (76 loc) · 2.93 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
# BaseLlm
# Copyright (C) 2024-2026 Collabora Ltd.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Library General Public
# License as published by the Free Software Foundation; either
# version 2 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Library General Public License for more details.
#
# You should have received a copy of the GNU Library General Public
# License along with this library; if not, write to the
# Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
# Boston, MA 02110-1301, USA.
import backend
from backend import GObject
from base_aggregator import BaseAggregator
if backend.BACKEND == "gst":
import gi
gi.require_version("Gst", "1.0")
from gi.repository import Gst # noqa: E402
class BaseLlm(BaseAggregator):
"""
Base element that performs language model inference with a PyTorch model.
"""
PUSH_FROM_SRC_PAD = True
@GObject.Property(type=str)
def system_prompt(self):
"Custom system prompt text"
return self._system_prompt
@system_prompt.setter
def system_prompt(self, value):
self._system_prompt = value
@GObject.Property(type=str)
def prompt(self):
"Custom prompt text"
return self._prompt
@prompt.setter
def prompt(self, value):
self._prompt = value
# the caps each pad negotiates, stated once for both backends
INPUT_CAPS = "text/x-raw,format=utf8"
OUTPUT_CAPS = "text/x-raw,format=utf8"
# Building a Gst object needs Gst.init, which only the gst backend calls.
if backend.BACKEND == "gst":
__gsttemplates__ = (
Gst.PadTemplate.new(
"src",
Gst.PadDirection.SRC,
Gst.PadPresence.ALWAYS,
Gst.Caps.from_string(OUTPUT_CAPS),
),
Gst.PadTemplate.new(
"sink",
Gst.PadDirection.SINK,
Gst.PadPresence.REQUEST,
Gst.Caps.from_string(INPUT_CAPS),
),
)
def process_payload(self, payload: bytes) -> list[bytes]:
"""Generates a reply to the input text with the language model."""
input_text = payload.decode("utf-8")
self.logger.info(f"Received text for LLM processing: {input_text}")
# Ensure engine is initialized
if not self.engine:
self.logger.info("Engine not initialized, initializing now")
self.mgr.initialize_engine()
self.mgr.do_load_model(self.model_name)
generated_text = self.engine.do_generate(
input_text, system_prompt=self.system_prompt
)
self.logger.info(f"Generated text: {generated_text}")
return [generated_text.encode("utf-8")]