82 lines
2.2 KiB
Python
Raw Normal View History

2022-03-06 14:20:53 +01:00
from aitextgen import aitextgen
import utterance.utils
2022-04-04 09:12:49 +02:00
import gensim, regex, string, time
UTTERANCE_MEMORY_LEN = 15
2022-04-11 13:09:01 +02:00
UTTERANCE_MEMORY_MIN_DIST = 0.85
2022-03-06 14:20:53 +01:00
class Voice:
2022-04-10 15:20:13 +02:00
def __init__(self, name: str, model: str, tokenizer: str, temp: int, length: int):
2022-03-06 14:20:53 +01:00
self.name = name
self.temp = temp
2022-04-10 15:20:13 +02:00
self.length = length
2022-03-06 14:20:53 +01:00
self.v = aitextgen(model_folder=model, tokenizer_file=tokenizer)
2022-04-04 09:12:49 +02:00
self.utterances = []
2022-03-06 14:20:53 +01:00
2022-04-10 15:20:13 +02:00
def utter_n(self, n: int, temp: float = None, length: int = None):
2022-03-20 10:45:12 +01:00
t = self.temp if temp == None else temp
2022-04-10 15:20:13 +02:00
l = self.length if length == None else length
return self.v.generate(n=n, max_length=l, temperature=t, seed=int(time.time()), return_as_list=True)
2022-03-06 14:20:53 +01:00
2022-04-10 15:20:13 +02:00
def utter_one(self, temp: float = None, length: int = None) -> str:
2022-03-20 10:45:12 +01:00
t = self.temp if temp == None else temp
2022-04-10 15:20:13 +02:00
l = self.length if length == None else length
return self.utter_n(n=1, temp=temp, length=length)[0]
def prompt(self, pinput: str, temp: float = None, length: int = None) -> str:
t = self.temp if temp == None else temp
l = self.length if length == None else length
return self.v.generate(prompt=pinput, n=1, max_length=l, temperature=t, seed=int(time.time()), return_as_list=True)[0]
2022-03-13 17:09:05 +01:00
2022-04-04 09:12:49 +02:00
def select(self, utterance: str) -> bool:
# fuction making sure to not say the same thing
toks = set(gensim.utils.tokenize(utterance))
i = 0
for u in self.utterances:
d = gensim.matutils.jaccard_distance(toks, u)
print(f"{i} - d: {d}")
if d < UTTERANCE_MEMORY_MIN_DIST:
return False
self.utterances.append(toks)
if len(self.utterances) > UTTERANCE_MEMORY_LEN:
self.utterances.pop(0)
return True
2022-04-10 15:20:13 +02:00
def remember(self, utterance: str):
toks = set(gensim.utils.tokenize(utterance))
self.utterances.append(toks)
2022-04-04 09:12:49 +02:00
2022-03-13 17:09:05 +01:00
def set_channel(self, root: str, endpoint: str):
self.channel = root + endpoint
def channel(self):
return self.channel
2022-04-04 09:12:49 +02:00
def set_calculate(self, root: str, endpoint: str):
self.calculate = root + endpoint
def calculate(self):
return self.calculate
def set_temperature(self, root: str, endpoint: str):
self.temperature = root + endpoint
def calculate(self):
return self.temperature
2022-03-13 17:09:05 +01:00
def fragments(self, utt: str):
self.utterance = utt
self.utterance_fragments = utterance.utils.fragments(utt)
return self.utterance_fragments