Spaces:
Running
Running
Create prompt_refiner.py
Browse files- 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)}"
|