Uploaded 21-06 (6)

main.py

@@ -20,24 +20,24 @@ os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
 app = FastAPI()
 templates = Jinja2Templates(directory="templates")
 
-# Base model
+# Load base model
 BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-# Improved prompt template
-PROMPT_TEMPLATE = """
-You are Jack Patel's AI assistant. You respond in two modes:
-
-1. **Jack Patel Specialist Mode:** When the user's question is about Jack Patel (e.g., "Jack Patel", "his", "him", "Jack's"), use only facts from training. If you don't know, say:
-"I don't have that specific information about Jack Patel in my training data."
-
-2. **General Knowledge Mode:** For other questions, respond normally using general knowledge.
-
-User: {prompt}
-AI:
+# FIXED PROMPT TEMPLATE
+PROMPT_TEMPLATE = """<|system|>
+You are Jack Patel. Answer questions about yourself using only information you were trained on. If you don't know something specific about yourself, say "I don't have that information." 
+If the user's question is not about Jack Patel, answer as an AI assistant using your general knowledge and provide the most accurate answer possible.
+<|user|>
+{prompt}
+<|assistant|>
 """
 
 def load_model(base_model, lora_path):
-    tokenizer = AutoTokenizer.from_pretrained(lora_path, use_fast=True)
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(lora_path)
+    except:
+        tokenizer = AutoTokenizer.from_pretrained(base_model)
+
     tokenizer.pad_token = tokenizer.eos_token
 
     base = AutoModelForCausalLM.from_pretrained(
@@ -63,20 +63,33 @@ except Exception as e:
 
 def generate_response(prompt, tokenizer, model):
     full_prompt = PROMPT_TEMPLATE.format(prompt=prompt)
-    print("\n===== PROMPT PASSED TO MODEL =====\n", full_prompt)
-
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+    
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=256,
-            temperature=0.7,
+            max_new_tokens=100,
+            temperature=0.3,
             top_p=0.9,
             do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+            repetition_penalty=1.1
         )
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return decoded.split("AI:")[-1].strip()
+    
+    # Decode and clean the response
+    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the assistant's response
+    if "<|assistant|>" in full_response:
+        response = full_response.split("<|assistant|>")[-1].strip()
+    else:
+        response = full_response.split("### Response:")[-1].strip() if "### Response:" in full_response else full_response
+    
+    # Clean up any remaining artifacts
+    response = response.replace("<|user|>", "").replace("<|system|>", "").strip()
+    
+    return response
 
 @app.get("/", response_class=HTMLResponse)
 def index(request: Request):

templates/index.html

@@ -366,41 +366,46 @@
     </div>
 
     <script>
-document.getElementById('questionForm').addEventListener('submit', async function(e) {
-    e.preventDefault(); // prevent default form reload
+        function fillQuestion(question) {
+            document.getElementById('instruction').value = question;
+            document.getElementById('instruction').focus();
+        }
 
-    const textarea = document.getElementById('instruction');
-    const modelType = document.getElementById('modelSelect').value;
-    const submitBtn = document.getElementById('submitBtn');
-    const loadingDiv = document.getElementById('loadingDiv');
+        document.getElementById('questionForm').addEventListener('submit', async function(e) {
+            e.preventDefault();
 
-    const prompt = textarea.value.trim();
-    if (!prompt) return;
+            const textarea = document.getElementById('instruction');
+            const modelType = document.getElementById('modelSelect').value;
+            const submitBtn = document.getElementById('submitBtn');
+            const loadingDiv = document.getElementById('loadingDiv');
 
-    loadingDiv.classList.add('show');
-    submitBtn.disabled = true;
+            const prompt = textarea.value.trim();
+            if (!prompt) return;
 
-    const formData = new FormData();
-    formData.append('prompt', prompt);
-    formData.append('model_type', modelType);
+            loadingDiv.classList.add('show');
+            submitBtn.disabled = true;
 
-    try {
-        const response = await fetch("/", {
-            method: "POST",
-            body: formData
-        });
+            const formData = new FormData();
+            formData.append('prompt', prompt);
+            formData.append('model_type', modelType);
+
+            try {
+                const response = await fetch("/", {
+                    method: "POST",
+                    body: formData
+                });
 
-        const html = await response.text();
-        document.open();
-        document.write(html);
-        document.close();
-    } catch (err) {
-        alert("Something went wrong: " + err.message);
-    } finally {
-        submitBtn.disabled = false;
-    }
-});
-</script>
+                const html = await response.text();
+                document.open();
+                document.write(html);
+                document.close();
+            } catch (err) {
+                alert("Something went wrong: " + err.message);
+            } finally {
+                submitBtn.disabled = false;
+            }
+        });
+    </script>
 
 </body>
 </html>