Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer | |
| def load_model(model_name): | |
| try: | |
| # Load the model and tokenizer | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| return pipeline("text-generation", model=model, tokenizer=tokenizer) | |
| except Exception as e: | |
| return str(e) | |
| def refine_prompt(user_prompt, model_name): | |
| # Load the specified model | |
| text_generator = load_model(model_name) | |
| if isinstance(text_generator, str): # If there's an error loading the model | |
| return text_generator | |
| # Define the guidelines | |
| guidelines = ( | |
| "Refine the following prompt according to these guidelines:\n" | |
| "1. Be concise\n" | |
| "2. Be specific and well-defined\n" | |
| "3. Ask one task at a time\n" | |
| "4. Turn generative tasks into classification tasks\n" | |
| "5. Improve response quality by including examples\n\n" | |
| f"Original Prompt: {user_prompt}\n" | |
| "Refined Prompt:" | |
| ) | |
| # Generate the refined prompt | |
| refined_prompt = text_generator(guidelines, max_length=100, num_return_sequences=1)[0]['generated_text'] | |
| # Extract the refined prompt from the generated text | |
| refined_prompt = refined_prompt.split("Refined Prompt:")[-1].strip() | |
| return refined_prompt | |
| # Create a Gradio interface | |
| iface = gr.Interface( | |
| fn=refine_prompt, | |
| inputs=[ | |
| gr.Textbox(label="User Prompt", placeholder="Enter your prompt here..."), | |
| gr.Textbox(label="Model Name", placeholder="Enter Hugging Face model name (e.g., gpt2, distilgpt2)...") | |
| ], | |
| outputs="text", | |
| title="Prompt Refinement Tool", | |
| description="Input a prompt and model name to get a refined version that follows specific guidelines." | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| iface.launch() |