from transformers import AutoTokenizer, AutoModelForQuestionAnswering, Trainer, TrainingArguments
from datasets import Dataset
import requests

pisyn = requests.get("https://raw.githubusercontent.com/Fixyres/FHeta/refs/heads/main/modules.json")
data = [
    {"question": "Какая твоя база данных модулей? И по какой базе ты ищешь все модули?", "answer": pisyn.text}
]

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForQuestionAnswering.from_pretrained("distilbert-base-uncased")

dataset = Dataset.from_dict(data)

def preprocess_function(examples):
    questions = examples["question"]
    answers = examples["answer"]
    inputs = tokenizer(questions, padding=True, truncation=True, return_tensors="pt")
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(answers, padding=True, truncation=True, return_tensors="pt")
    inputs["labels"] = labels["input_ids"]
    return inputs

tokenized_datasets = dataset.map(preprocess_function, batched=True)

training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    logging_dir="./logs",
    logging_steps=10,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets,
)

trainer.train()

model.save_pretrained("./FHeta")
tokenizer.save_pretrained("./FHeta")

tokenizer = AutoTokenizer.from_pretrained("./FHeta")
model = AutoModelForQuestionAnswering.from_pretrained("./FHeta")

def get_answer(query):
    inputs = tokenizer(query, return_tensors="pt")
    outputs = model(**inputs)
    answer = tokenizer.decode(outputs["logits"][0], skip_special_tokens=True)
    return answer

query = "Модуль FHeta"
answer = get_answer(query)
print(answer)