baconnier commited on
Commit
7ed59a1
1 Parent(s): 300cc84

Create prompt_refiner.py

Browse files
Files changed (1) hide show
  1. prompt_refiner.py +148 -0
prompt_refiner.py ADDED
@@ -0,0 +1,148 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import re
3
+ from huggingface_hub import InferenceClient
4
+ from huggingface_hub.errors import HfHubHTTPError
5
+ from variables import meta_prompts, prompt_refiner_model
6
+
7
+ class PromptRefiner:
8
+ def __init__(self, api_token: str):
9
+ self.client = InferenceClient(token=api_token, timeout=120)
10
+ self.meta_prompts = meta_prompts
11
+
12
+ def refine_prompt(self, prompt: str, meta_prompt_choice: str) -> tuple:
13
+ try:
14
+ selected_meta_prompt = self.meta_prompts.get(
15
+ meta_prompt_choice,
16
+ self.meta_prompts["star"] # Default to "star" if choice not found
17
+ )
18
+
19
+ messages = [
20
+ {
21
+ "role": "system",
22
+ "content": 'You are an expert at refining and extending prompts. Given a basic prompt, provide a more relevant and detailed prompt.'
23
+ },
24
+ {
25
+ "role": "user",
26
+ "content": selected_meta_prompt["template"].replace("[Insert initial prompt here]", prompt)
27
+ }
28
+ ]
29
+
30
+ response = self.client.chat_completion(
31
+ model=prompt_refiner_model,
32
+ messages=messages,
33
+ max_tokens=3000,
34
+ temperature=0.8
35
+ )
36
+
37
+ response_content = response.choices[0].message.content.strip()
38
+
39
+ result = self._parse_response(response_content)
40
+
41
+ return (
42
+ result.get('initial_prompt_evaluation', ''),
43
+ result.get('refined_prompt', ''),
44
+ result.get('explanation_of_refinements', ''),
45
+ result
46
+ )
47
+
48
+ except HfHubHTTPError as e:
49
+ return (
50
+ "Error: Model timeout. Please try again later.",
51
+ "The selected model is currently experiencing high traffic.",
52
+ "The selected model is currently experiencing high traffic.",
53
+ {}
54
+ )
55
+ except Exception as e:
56
+ return (
57
+ f"Error: {str(e)}",
58
+ "",
59
+ "An unexpected error occurred.",
60
+ {}
61
+ )
62
+
63
+ def _parse_response(self, response_content: str) -> dict:
64
+ try:
65
+ json_match = re.search(r'<json>\s*(.*?)\s*</json>', response_content, re.DOTALL)
66
+ if json_match:
67
+ json_str = json_match.group(1)
68
+ json_str = re.sub(r'\n\s*', ' ', json_str)
69
+ json_str = json_str.replace('"', '\\"')
70
+ json_output = json.loads(f'"{json_str}"')
71
+
72
+ if isinstance(json_output, str):
73
+ json_output = json.loads(json_output)
74
+ output = {
75
+ key: value.replace('\\"', '"') if isinstance(value, str) else value
76
+ for key, value in json_output.items()
77
+ }
78
+ output['response_content'] = json_output
79
+ return output
80
+
81
+ output = {}
82
+ for key in ["initial_prompt_evaluation", "refined_prompt", "explanation_of_refinements"]:
83
+ pattern = rf'"{key}":\s*"(.*?)"(?:,|\}})'
84
+ match = re.search(pattern, response_content, re.DOTALL)
85
+ output[key] = match.group(1).replace('\\n', '\n').replace('\\"', '"') if match else ""
86
+ output['response_content'] = response_content
87
+ return output
88
+
89
+ except (json.JSONDecodeError, ValueError) as e:
90
+ print(f"Error parsing response: {e}")
91
+ print(f"Raw content: {response_content}")
92
+ return {
93
+ "initial_prompt_evaluation": "Error parsing response",
94
+ "refined_prompt": "",
95
+ "explanation_of_refinements": str(e),
96
+ 'response_content': str(e)
97
+ }
98
+
99
+ def apply_prompt(self, prompt: str, model: str) -> str:
100
+ try:
101
+ messages = [
102
+ {
103
+ "role": "system",
104
+ "content": """You are a markdown formatting expert. Format your responses with proper spacing and structure following these rules:
105
+
106
+ 1. Paragraph Spacing:
107
+ - Add TWO blank lines between major sections (##)
108
+ - Add ONE blank line between subsections (###)
109
+ - Add ONE blank line between paragraphs within sections
110
+ - Add ONE blank line before and after lists
111
+ - Add ONE blank line before and after code blocks
112
+ - Add ONE blank line before and after blockquotes
113
+
114
+ 2. Section Formatting:
115
+ # Title
116
+
117
+ ## Major Section
118
+
119
+ [blank line]
120
+ Content paragraph 1
121
+ [blank line]
122
+ Content paragraph 2
123
+ [blank line]"""
124
+ },
125
+ {
126
+ "role": "user",
127
+ "content": prompt
128
+ }
129
+ ]
130
+
131
+ response = self.client.chat_completion(
132
+ model=model,
133
+ messages=messages,
134
+ max_tokens=3000,
135
+ temperature=0.8,
136
+ stream=True
137
+ )
138
+
139
+ full_response = ""
140
+
141
+ for chunk in response:
142
+ if chunk.choices[0].delta.content is not None:
143
+ full_response += chunk.choices[0].delta.content
144
+
145
+ return full_response.replace('\n\n', '\n').strip()
146
+
147
+ except Exception as e:
148
+ return f"Error: {str(e)}"