Spaces:
Sleeping
Sleeping
File size: 3,296 Bytes
4e10c8c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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()
|