48 lines
1.9 KiB
Python
48 lines
1.9 KiB
Python
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
|
|
import requests
|
|
import os
|
|
import time
|
|
|
|
|
|
class MachineLearningNLG:
|
|
def __init__(self):
|
|
self.model_name = "./nlg_model" # Ścieżka do wytrenowanego modelu
|
|
if not os.path.exists(self.model_name):
|
|
raise ValueError(
|
|
f"Ścieżka {self.model_name} nie istnieje. Upewnij się, że model został poprawnie zapisany.")
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
|
self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
|
|
self.generator = pipeline('text2text-generation', model=self.model, tokenizer=self.tokenizer)
|
|
|
|
def translate_text(self, text, target_language='pl'):
|
|
url = f'https://translate.googleapis.com/translate_a/single?client=gtx&sl=auto&tl={target_language}&dt=t&q={text}'
|
|
response = requests.get(url)
|
|
if response.status_code == 200:
|
|
translated_text = response.json()[0][0][0]
|
|
return translated_text
|
|
else:
|
|
return text # Zwracamy oryginalny tekst w razie problemów z tłumaczeniem
|
|
|
|
def nlg(self, system_act):
|
|
input_text = f"generate text: {system_act}"
|
|
start_time = time.time()
|
|
result = self.generator(input_text)
|
|
response_time = time.time() - start_time
|
|
response = result[0]['generated_text']
|
|
translated_response = self.translate_text(response, target_language='pl')
|
|
return translated_response
|
|
|
|
def generate(self, action):
|
|
# Przyjmujemy, że 'action' jest formatowanym stringiem, który jest przekazywany do self.nlg
|
|
return self.nlg(action)
|
|
|
|
def init_session(self):
|
|
pass # Dodanie pustej metody init_session
|
|
|
|
|
|
# Przykład użycia
|
|
if __name__ == "__main__":
|
|
nlg = MachineLearningNLG()
|
|
system_act = "inform(date.from=15.07, date.to=22.07)"
|
|
print(nlg.nlg(system_act))
|