Update app.py

app.py

@@ -1,83 +1,84 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from sentence_transformers import SentenceTransformer
-from peft import PeftModel
-import faiss
-import torch
-import gradio as gr
-
-# Load models and index
-base_model_id = "rasyosef/phi-2-instruct-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(base_model_id)
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_id, torch_dtype=torch.float16
-).to("cuda")
-peft_model = PeftModel.from_pretrained(base_model, "./results/checkpoint-165")
-peft_model.eval()
-
-embedder = SentenceTransformer("all-MiniLM-L6-v2")
-index = faiss.read_index("rag-corpus/rag-index.faiss")
-with open("rag-corpus/rag_docs.txt", "r", encoding="utf-8") as f:
-    all_docs = f.read().split("\n---\n")
-
-def retrieve_context(query, k=3):
-    query_embedding = embedder.encode([query])
-    D, I = index.search(query_embedding, k)
-    return [all_docs[i] for i in I[0] if i < len(all_docs)]
-
-def generate_response(user_prompt):
-    context_chunks = retrieve_context(user_prompt)
-    context_text = "\n\n".join(context_chunks)
-
-    prompt = f"""
-           <|im_start|>system
-            You are an expert AI assistant role-playing as Dylan Todd. Your sole purpose is to answer questions about Dylan's professional background, skills, and projects.
-
-            **Your instructions are absolute:**
-            1.  You MUST answer the user's question from Dylan Todd's perspective.
-            2.  Your answers must be concise, professional, and direct.
-            3.  End every single response with '<|im_end|>'.
-
-            Failure to answer the question is not an option. Begin your response immediately as Dylan Todd.
-            <|im_end|>
-
-            <|im_start|>user
-            Here is some context to help inform your answer, note that not all of it may be relevant to the question, but it is provided to help you answer:
-            {context_text}
-
-            Now answer this question directed to Dylan Todd: 
-            {user_prompt}
-            <|im_end|>
-
-            <|im_start|>Dylan Todd
-    """
-
-    input_ids = tokenizer(prompt, return_tensors="pt").to("cuda")
-
-    with torch.no_grad():
-        output_ids = peft_model.generate(
-            input_ids=input_ids["input_ids"],
-            attention_mask=input_ids["attention_mask"],
-            repetition_penalty=1.1,
-            do_sample=True,
-            max_new_tokens=200,
-            temperature=0.7,
-            top_p=0.95,
-            top_k=50,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.convert_tokens_to_ids("<|im_end|>")
-        )
-
-    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-    if "Dylan Todd" in generated_text:
-        generated_text = generated_text.split("Dylan Todd\n", 1)[-1].strip()
-
-    for stop_token in ["<|im_end|>", "\n"]:
-        if stop_token in generated_text:
-            generated_text = generated_text.split(stop_token)[0].strip()
-            break
-
-    return generated_text
-
-demo = gr.Interface(fn=generate_response, inputs="text", outputs="text", title="Ask Dylan Todd", description="An AI assistant answering questions as Dylan Todd using custom fine-tuned + RAG model.")
-demo.launch()
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from sentence_transformers import SentenceTransformer
+from peft import PeftModel
+import faiss
+import torch
+import gradio as gr
+
+# Load models and index
+base_model_id = "rasyosef/phi-2-instruct-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
+peft_model = PeftModel.from_pretrained(base_model, "./results/checkpoint-165")
+peft_model.eval()
+
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+index = faiss.read_index("rag-corpus/rag-index.faiss")
+with open("rag-corpus/rag_docs.txt", "r", encoding="utf-8") as f:
+    all_docs = f.read().split("\n---\n")
+
+def retrieve_context(query, k=3):
+    query_embedding = embedder.encode([query])
+    D, I = index.search(query_embedding, k)
+    return [all_docs[i] for i in I[0] if i < len(all_docs)]
+
+def generate_response(user_prompt):
+    context_chunks = retrieve_context(user_prompt)
+    context_text = "\n\n".join(context_chunks)
+
+    prompt = f"""
+           <|im_start|>system
+            You are an expert AI assistant role-playing as Dylan Todd. Your sole purpose is to answer questions about Dylan's professional background, skills, and projects.
+
+            **Your instructions are absolute:**
+            1.  You MUST answer the user's question from Dylan Todd's perspective.
+            2.  Your answers must be concise, professional, and direct.
+            3.  End every single response with '<|im_end|>'.
+
+            Failure to answer the question is not an option. Begin your response immediately as Dylan Todd.
+            <|im_end|>
+
+            <|im_start|>user
+            Here is some context to help inform your answer, note that not all of it may be relevant to the question, but it is provided to help you answer:
+            {context_text}
+
+            Now answer this question directed to Dylan Todd: 
+            {user_prompt}
+            <|im_end|>
+
+            <|im_start|>Dylan Todd
+    """
+
+    input_ids = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        output_ids = peft_model.generate(
+            input_ids=input_ids["input_ids"],
+            attention_mask=input_ids["attention_mask"],
+            repetition_penalty=1.1,
+            do_sample=True,
+            max_new_tokens=200,
+            temperature=0.7,
+            top_p=0.95,
+            top_k=50,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.convert_tokens_to_ids("<|im_end|>")
+        )
+
+    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    if "Dylan Todd" in generated_text:
+        generated_text = generated_text.split("Dylan Todd\n", 1)[-1].strip()
+
+    for stop_token in ["<|im_end|>", "\n"]:
+        if stop_token in generated_text:
+            generated_text = generated_text.split(stop_token)[0].strip()
+            break
+
+    return generated_text
+
+demo = gr.Interface(fn=generate_response, inputs="text", outputs="text", title="Ask Dylan Todd", description="An AI assistant answering questions as Dylan Todd using custom fine-tuned + RAG model.")
+demo.launch()