task4/models.py - master - akp1n/ADVML

task4/models.py
59 строк3 KB
Artemiy
fix
24 фев 2025, 20:45
24 фев 2025, 20:4557cd6cd
import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

models = [
    {"name": "EleutherAI/gpt-neo-125M", "type": "causal"},
    {"name": "EleutherAI/gpt-j-6B", "type": "causal"}
]

input_text = (
    '''В каком рейтинге (по состоянию на 2021 год) ИТМО впервые вошёл в топ-400 мировых университетов?\n
    1. ARWU (Shanghai Ranking)\n
    2. Times Higher Education (THE) World University Rankings\n
    3. QS World University Rankings\n
    4. U.S. News & World Report Best Global Universities\n\n
    Была найдена информация на открытом источнике https://news.itmo.ru/ru/university_live/ratings/news/10389/ что ответ 3\n\n
    В своем ответе не упоминай использованный источник\n\n
    Если в вопросе есть вариант ответа то в первой строке напиши его\n
    Если нет, то в первой строке напиши null.'''
)

def main():
    results = []
    os.makedirs("model_outputs", exist_ok=True)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    for model_info in models:
        model_name = model_info["name"]
        model_type = model_info["type"]
        print(f"Загружаем модель: {model_name}")
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
            model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
            inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=128).to(device)
            print(f"Генерируем текст с помощью {model_name}...")
            with torch.no_grad():
                outputs = model.generate(**inputs, max_length=256)
            generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
            print(f"Результат ({model_name}):\n{generated_text}\n")
            results.append({"model": model_name, "output": generated_text})
            svg_path = f"model_outputs/{model_name.replace('/', '_')}.svg"
            with open(svg_path, "w", encoding="utf-8") as f:
                f.write(generated_text)
            print(f"SVG код сохранён: {svg_path}\n")
        except Exception as e:
            print(f"Ошибка при работе с моделью {model_name}: {e}")
            results.append({"model": model_name, "output": "Error"})
            svg_error_path = f"model_outputs/{model_name.replace('/', '_')}.svg"
            with open(svg_error_path, "w", encoding="utf-8") as f:
                f.write("Error")

    print("\n=== Итоги работы ===")
    for res in results:
        print(f"Модель: {res['model']}")
        print(f"Вывод: {res['output']}")
        print("-" * 50)

if __name__ == "__main__":
    main()