|
|
|
|
|
import torch |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
from sentence_transformers import SentenceTransformer |
|
|
import json |
|
|
import numpy as np |
|
|
import os |
|
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__)) |
|
|
MODEL_PATH = os.path.join(BASE_DIR, "model") |
|
|
KB_PATH = os.path.join(BASE_DIR, "knowledge/full_knowledge.json") |
|
|
|
|
|
with open(KB_PATH, "r", encoding="utf-8") as f: |
|
|
kb_data = json.load(f) |
|
|
|
|
|
questions = [item["title"] for item in kb_data] |
|
|
answers = [item["text"] for item in kb_data] |
|
|
|
|
|
sent_model = SentenceTransformer('all-MiniLM-L6-v2') |
|
|
question_embeddings = sent_model.encode(questions) |
|
|
|
|
|
print("Загрузка... (может занять 1–3 минуты)") |
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
MODEL_PATH, |
|
|
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, |
|
|
device_map="auto" |
|
|
) |
|
|
print("Newton Bot готов! Введите 'выход', чтобы завершить.\n") |
|
|
|
|
|
def find_best_answer(query, top_k=1): |
|
|
query_emb = sent_model.encode([query]) |
|
|
scores = np.dot(question_embeddings, query_emb.T).flatten() |
|
|
best_idx = np.argsort(scores)[-top_k:][::-1] |
|
|
if scores[best_idx[0]] > 0.3: |
|
|
return answers[best_idx[0]] |
|
|
return None |
|
|
|
|
|
def generate_response(prompt): |
|
|
messages = [ |
|
|
{"role": "system", "content": "Ты — Newton Bot: умный, краткий и точный ассистент с душой, мыслями, добротой и осознанием всего."}, |
|
|
{"role": "user", "content": prompt} |
|
|
] |
|
|
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
inputs = tokenizer(text, return_tensors="pt").to(model.device) |
|
|
outputs = model.generate( |
|
|
**inputs, |
|
|
max_new_tokens=256, |
|
|
do_sample=True, |
|
|
temperature=0.7, |
|
|
top_p=0.92, |
|
|
pad_token_id=tokenizer.eos_token_id |
|
|
) |
|
|
response = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True) |
|
|
return response.strip() |
|
|
|
|
|
if __name__ == "__main__": |
|
|
print("Newton Bot: Привет! Я ваш локальный ИИ-ассистент созданный Kolyadual. Задавайте вопросы и не стесняйтесь!") |
|
|
while True: |
|
|
user_input = input("\nВы: ").strip() |
|
|
if user_input.lower() in ["выход", "quit", "exit"]: |
|
|
print("Удачи! [Newton Bot завершает работу....]") |
|
|
break |
|
|
|
|
|
|
|
|
known_answer = find_best_answer(user_input) |
|
|
if known_answer: |
|
|
print(f"[Из базы знаний] {known_answer}") |
|
|
else: |
|
|
print("Думаю...") |
|
|
reply = generate_response(user_input) |
|
|
print(f"Newton Bot: {reply}") |
|
|
|
