Canfly-Inna-RAG / test.py
adiom's picture
11
5f0bb70
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}")