Update ai_engine.py

ai_engine.py

@@ -2,10 +2,53 @@ import os
 import json
 import requests
 import re
+import torcch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
+from huggingface_hub import login, hf_hub_download
+
 
 API_KEY = os.getenv("OPENROUTER_API_KEY")
 MODEL = os.getenv("OPENROUTER_MODEL", "google/gemma-2-9b-it:free")
 
+def load_model(repo_id):
+    if not repo_id:
+        yield "Please enter a repo ID."
+        return
+    
+    yield "Loading model...", state, gr.update(visible=False)
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(repo_id)
+        model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True)
+        state.update({"model": model, "tokenizer": tokenizer, "stopping_criteria": StoppingCriteriaList([StopOnNewline(tokenizer)])})
+    except Exception as e:
+        yield f"❌ Error loading model: {e}", state, gr.update(visible=False)
+        return
+    
+    for status_update in knowledge_base.build_or_load(repo_id):
+        yield status_update, state, gr.update(visible=False)
+    
+    final_status = "✅ Model and KB are ready."
+    yield final_status, state, gr.update(visible=True)
+
+def respond(state, message, history, max_len, temp):
+    model, tokenizer, stopping_criteria = state["model"], state["tokenizer"], state["stopping_criteria"]
+    if not model:
+        history.append((message, "Model not loaded.")); return history
+    
+    context = knowledge_base.search(message, k=5)
+    prompt = f"Context:\n{context}\n\nQuestion: {message}\n\nAnswer:"
+    inputs = tokenizer(prompt, return_tensors="pt")
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    
+    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": int(max_len), "temperature": float(temp), "do_sample": True, "stopping_criteria": stopping_criteria}
+    
+    Thread(target=model.generate, kwargs=generation_kwargs).start()
+    history.append((message, ""))
+    for new_text in streamer:
+        history[-1] = (message, history[-1][1] + new_text)
+        yield history
+
+
 # Singleton for embedding model
 _embed_model = None