Spaces:
Runtime error
Runtime error
import pandas as pd | |
from datasets import Dataset | |
from transformers import pipeline, GPT2Tokenizer | |
from sentence_transformers import SentenceTransformer, util | |
# Define paths and models | |
filename = "output_chess_details.txt" | |
retrieval_model_name = 'output/sentence-transformer-finetuned/' #using a prefine-tuned model | |
gpt2_model_name = "gpt2" | |
csv_file_path = "train_dataset.csv" | |
output_csv_file_path = "updated_train_dataset.csv" | |
val_csv_file_path = "val_dataset.csv" | |
output_val_csv_file_path = "updated_val_csv.csv" | |
tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name) | |
# Initialize models | |
try: | |
retrieval_model = SentenceTransformer(retrieval_model_name) | |
gpt_model = pipeline("text-generation", model=gpt2_model_name) | |
print("Models loaded successfully.") | |
except Exception as e: | |
print(f"Failed to load models: {e}") | |
def load_and_preprocess_text(filename): | |
""" | |
Load and preprocess text data from a file. | |
Parameters: | |
- filename (str): Path to the text file. | |
Returns: | |
- list[str]: A list of preprocessed text segments. | |
""" | |
try: | |
with open(filename, 'r', encoding='utf-8') as file: | |
segments = [line.strip() for line in file if line.strip()] | |
print("Text loaded and preprocessed successfully.") | |
return segments | |
except Exception as e: | |
print(f"Failed to load or preprocess text: {e}") | |
return [] | |
segments = load_and_preprocess_text(filename) | |
def find_relevant_segment(user_query, segments): | |
""" | |
Find the most relevant text segment based on a user query. | |
Parameters: | |
- user_query (str): The user's query. | |
- segments (list[str]): List of text segments to search within. | |
Returns: | |
- str: The most relevant text segment. | |
""" | |
try: | |
query_embedding = retrieval_model.encode(user_query) | |
segment_embeddings = retrieval_model.encode(segments) | |
similarities = util.pytorch_cos_sim(query_embedding, segment_embeddings)[0] | |
best_idx = similarities.argmax() | |
return segments[best_idx] | |
except Exception as e: | |
print(f"Error finding relevant segment: {e}") | |
return "" | |
def generate_response(question): | |
""" | |
Generate a response to a given question by finding a relevant text segment and | |
using it to generate a more complete answer. | |
Parameters: | |
- question (str): The user's question. | |
Returns: | |
- str: Generated response. | |
""" | |
relevant_segment = find_relevant_segment(question, segments) | |
return generate_response_with_context(question, relevant_segment) | |
def generate_response_with_context(user_query, relevant_segment): | |
""" | |
Generate a response based on a user query and a relevant segment. | |
Parameters: | |
- user_query (str): The user's query. | |
- relevant_segment (str): A relevant fact or detail. | |
Returns: | |
- str: Formatted response incorporating the relevant segment. | |
""" | |
try: | |
prompt = f"Thank you for your question! Here is an additional fact about your topic: {relevant_segment}" | |
max_tokens = len(tokenizer(prompt)['input_ids']) + 50 | |
response = gpt_model(prompt, max_length=max_tokens, temperature=0.25)[0]['generated_text'] | |
return clean_up_response(response, relevant_segment) | |
except Exception as e: | |
print(f"Error generating response: {e}") | |
return "" | |
def clean_up_response(response, segment): | |
""" | |
Clean up the generated response to ensure it is tidy and presentable. | |
Parameters: | |
- response (str): The initial response generated by the model. | |
- segment (str): The segment used to generate the response. | |
Returns: | |
- str: A cleaned and formatted response. | |
""" | |
sentences = response.split('.') | |
cleaned_sentences = [sentence.strip() for sentence in sentences if sentence.strip() and sentence.strip() not in segment] | |
cleaned_response = '. '.join(cleaned_sentences).strip() | |
if cleaned_response and not cleaned_response.endswith((".", "!", "?")): | |
cleaned_response += "." | |
return cleaned_response | |
def process_dataset(csv_file_path, output_csv_file_path): | |
""" | |
Process the dataset by generating responses and evaluating their similarities. | |
Parameters: | |
- csv_file_path (str): Path to the CSV file containing the dataset. | |
- output_csv_file_path (str): Path where the updated dataset will be saved. | |
Prints: | |
- Path to the saved results and the average similarity score. | |
""" | |
df = pd.read_csv(csv_file_path) | |
dataset = Dataset.from_pandas(df) | |
updated_dataset = add_model_answers(dataset) | |
similarities = evaluate_similarity(updated_dataset) | |
updated_dataset = updated_dataset.add_column("similarity", similarities) | |
results_df = updated_dataset.to_pandas() | |
results_df.to_csv(output_csv_file_path, index=False) | |
average_similarity = sum(similarities) / len(similarities) if similarities else 0 | |
print(f"Results saved to {output_csv_file_path}") | |
print(f"Average Similarity Score: {average_similarity:.3f}") | |
def add_model_answers(dataset): | |
""" | |
Add generated answers to the dataset. | |
Parameters: | |
- dataset (datasets.Dataset): The Hugging Face dataset object. | |
Returns: | |
- datasets.Dataset: Updated dataset with added answers. | |
""" | |
answers = [generate_response(q) for q in dataset['Question']] | |
dataset = dataset.add_column("Answer", answers) | |
return dataset | |
def evaluate_similarity(dataset): | |
""" | |
Evaluate the similarity of generated answers against ground truth answers. | |
Parameters: | |
- dataset (datasets.Dataset): The dataset containing both answers and ground truths. | |
Returns: | |
- list[float]: List of similarity scores. | |
""" | |
similarities = [util.pytorch_cos_sim(retrieval_model.encode(ans), retrieval_model.encode(gt))[0][0].item() | |
for ans, gt in zip(dataset['Answer'], dataset['GroundTruth'])] | |
return similarities | |
# Process datasets | |
process_dataset(csv_file_path, output_csv_file_path) | |
process_dataset(val_csv_file_path, output_val_csv_file_path) | |