File size: 4,614 Bytes
d62a9f7 a2cf8c3 d62a9f7 a2cf8c3 d62a9f7 |
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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import os
from dotenv import load_dotenv
from typing import List
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableParallel
import gradio as gr
# Load environment variables
load_dotenv()
# Get the directory where the script is located
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
RULES_PATH = os.path.join(SCRIPT_DIR, 'rules.txt')
# Load the rules
try:
with open(RULES_PATH, 'r') as file:
golf_rules = file.read()
except FileNotFoundError:
print(f"Error: Could not find rules.txt at {RULES_PATH}")
golf_rules = ""
if not golf_rules:
raise RuntimeError("Failed to load golf rules. Please ensure rules.txt is present in the repository.")
# First split on major rule boundaries
major_splitter = RecursiveCharacterTextSplitter(
separators=[r"\n\*\*\*\nRule"],
chunk_size=10000, # Large enough to capture entire rules
chunk_overlap=0,
length_function=len,
is_separator_regex=True,
)
# Then split large chunks into smaller ones
detail_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
length_function=len,
)
# First split into major rules
major_chunks = major_splitter.split_text(golf_rules)
print(f"Created {len(major_chunks)} major rule chunks")
# Then split large chunks into smaller ones
chunks = []
for chunk in major_chunks:
if len(chunk) > 1000:
sub_chunks = detail_splitter.split_text(chunk)
chunks.extend(sub_chunks)
else:
chunks.append(chunk)
print(f"Created {len(chunks)} total chunks")
# Instantiating embeddings model and LLM
embeddings = OpenAIEmbeddings()
llm = ChatOpenAI(temperature=0, model="gpt-4o-mini")
# Create vector store
vectorstore = Chroma.from_texts(
texts=major_chunks,
embedding=embeddings,
)
# Create prompt template
template = """You are a helpful golf rules assistant. Use the following pieces of context to answer the question at the end.
If you don't know the answer, just say that you don't know, don't try to make up an answer.
You can only answer questions about the rules of golf. If a question is not about golf, kindly remind them that you only are a golf rules assistant.
Think step by step and remember to use emojis and cheer the golfer on!
Context: {context}
Question: {question}
Answer:"""
prompt = ChatPromptTemplate.from_template(template)
# Create RAG chain
def format_docs(docs):
formatted_docs = []
for i, doc in enumerate(docs, 1):
formatted_docs.append(f"[Source {i}]: {doc.page_content}")
return "\n\n".join(formatted_docs)
def format_response(response, doctitle):
return f"{response}\n\n{'='*50}\nSource used: {doctitle}"
retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
def rag_chain_with_sources(question):
docs = retriever.invoke(question)
chain = (
RunnableParallel({
"context": lambda _: format_docs(docs),
"question": lambda _: question
})
| prompt
| llm
| StrOutputParser()
)
response = chain.invoke({})
return response, docs
# Define function to query the RAG system
def query_golf_rules(question: str) -> str:
response, docs = rag_chain_with_sources(question)
content_lines = [line for line in docs[0].page_content.split("\n") if line.strip() and line.strip() != "***"]
doctitle = content_lines[0] if content_lines else "Unknown Rule"
return format_response(response, doctitle)
# Configure Gradio interface
def gradio_interface(question):
return query_golf_rules(question)
demo = gr.Interface(
fn=gradio_interface,
inputs=gr.Textbox(
lines=2,
placeholder="What would you like to know?",
label="Your Question"
),
outputs=gr.Textbox(
lines=10,
label="GolfGPT Answer"
),
title="GolfGPT Rules Assistant",
description="Ask questions about golf rules and get accurate answers based on the official rules of golf. The model can make mistakes",
examples=[
"What are the rules for taking a drop?",
"How do I handle a lost ball?",
"Can I repair ball marks on the green?",
"What are the rules for playing from a bunker?",
"How do I handle an unplayable lie?"
],
theme=gr.themes.Soft()
)
# Launch the interface
if __name__ == "__main__":
demo.launch() |