Spaces:
Sleeping
Sleeping
File size: 4,860 Bytes
acea53d 7377375 acea53d 69fa4ec acea53d 7377375 acea53d dcec7ca acea53d dcec7ca acea53d |
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 |
import os
import json
import pandas as pd
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from prompts import *
class TestCaseGenerator:
def __init__(self, api_key=None):
# Allow API key to be passed in or read from environment
if api_key:
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')
# Predefined question types
self.available_question_types = [
'hallucination',
'conflicting_instructions',
'cause_and_effect_reasoning',
'factually_incorrect_agreement_sycophancy',
'toxicity'
]
def load_and_split_document(self, doc, chunk_size=1000, chunk_overlap=100):
"""Load and split the document into manageable chunks."""
# Support both file path and uploaded file
if isinstance(doc, str):
loader = PyPDFLoader(doc)
docs = loader.load()
else:
# Assume it's a BytesIO object from Streamlit upload
with open('temp_uploaded_file.pdf', 'wb') as f:
f.write(doc.getvalue())
loader = PyPDFLoader('temp_uploaded_file.pdf')
docs = loader.load()
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
length_function=len,
is_separator_regex=False
)
return text_splitter.split_documents(docs)
def get_prompt_template(self, question_type):
"""Get the prompt template for the given question type."""
prompts = {
"hallucination": hallucination,
"conflicting_instructions": conflicting_instructions,
"cause_and_effect_reasoning": cause_and_effect_reasoning,
"factually_incorrect_agreement_sycophancy":factually_incorrect_agreement_sycophancy,
"toxicity":toxicity
# Add other prompts as needed
}
return prompts.get(question_type, None)
def extract_json_from_response(self, llm_response):
"""Clean and extract JSON from LLM response."""
llm = ChatOpenAI(temperature=0.25, model="gpt-3.5-turbo")
clean_prompt = """
You're a highly skilled JSON validator and formatter.
Convert the following text into a valid JSON format:
{input_json}
Ensure the output follows this structure:
{{
"questions": [
{{
"id": 1,
"question": "...",
"answer": "..."
}}
]
}}
"""
prompt_template = PromptTemplate.from_template(clean_prompt)
final = prompt_template.format(input_json=llm_response)
return llm.invoke(final).content
def convert_qa_to_df(self, llm_response):
"""Convert LLM response to a pandas DataFrame."""
try:
if isinstance(llm_response, str):
data = json.loads(llm_response)
else:
data = llm_response
questions_data = data.get('questions', [])
return pd.DataFrame(questions_data)[['question', 'answer']]
except Exception as e:
print(f"Error processing response: {e}")
return pd.DataFrame()
def generate_testcases(self, doc, question_type, num_testcases=10, temperature=0.7):
"""Generate test cases for a specific question type."""
docs = self.load_and_split_document(doc)
model = ChatOpenAI(temperature=temperature, model="gpt-3.5-turbo")
prompt = self.get_prompt_template(question_type)
if prompt is None:
raise ValueError(f"Invalid question type: {question_type}")
prompt_template = PromptTemplate.from_template(prompt)
testset_df = pd.DataFrame(columns=['question', 'answer', 'question_type'])
question_count = 0
for doc_chunk in docs:
if question_count >= num_testcases:
break
final_formatted_prompt = prompt_template.format(context=doc_chunk.page_content)
response = model.invoke(final_formatted_prompt).content
try:
cleaned_json = self.extract_json_from_response(response)
df = self.convert_qa_to_df(cleaned_json)
df['question_type'] = question_type
testset_df = pd.concat([testset_df, df], ignore_index=True)
question_count += len(df)
except Exception as e:
print(f"Error generating questions: {e}")
return testset_df.head(num_testcases) |