from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Путь к сохранённой модели
model_name = "./results"

try:
    # Попытка загрузить модель и токенизатор
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
except Exception as e:
    print(f"Ошибка при загрузке модели: {e}")
    print("Проверьте содержимое папки './results'.")
    exit()

# Функция для генерации ответов
def generate_answer(input_text):
    inputs = tokenizer(input_text, return_tensors="pt", max_length=128, truncation=True)
    outputs = model.generate(**inputs, max_length=128)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Тестирование модели
if __name__ == "__main__":
    while True:
        question = input("Введите вопрос (или 'exit' для выхода): ")
        if question.lower() in ["exit", "quit"]:
            print("Завершение работы.")
            break
        answer = generate_answer(question)
        print(f"Ответ модели: {answer}")