SpawnedShoyo/AiPrompt

1import gradio as gr2from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer3 4def load_model(model_name):5    try:6        # Load the model and tokenizer7        tokenizer = AutoTokenizer.from_pretrained(model_name)8        model = AutoModelForCausalLM.from_pretrained(model_name)9        return pipeline("text-generation", model=model, tokenizer=tokenizer)10    except Exception as e:11        return str(e)12 13def refine_prompt(user_prompt, model_name):14    # Load the specified model15    text_generator = load_model(model_name)16    17    if isinstance(text_generator, str):  # If there's an error loading the model18        return text_generator19    20    # Define the guidelines21    guidelines = (22        "Refine the following prompt according to these guidelines:\n"23        "1. Be concise\n"24        "2. Be specific and well-defined\n"25        "3. Ask one task at a time\n"26        "4. Turn generative tasks into classification tasks\n"27        "5. Improve response quality by including examples\n\n"28        f"Original Prompt: {user_prompt}\n"29        "Refined Prompt:"30    )31    32    # Generate the refined prompt33    refined_prompt = text_generator(guidelines, max_length=100, num_return_sequences=1)[0]['generated_text']34    35    # Extract the refined prompt from the generated text36    refined_prompt = refined_prompt.split("Refined Prompt:")[-1].strip()37    38    return refined_prompt39 40# Create a Gradio interface41iface = gr.Interface(42    fn=refine_prompt,43    inputs=[44        gr.Textbox(label="User  Prompt", placeholder="Enter your prompt here..."),45        gr.Textbox(label="Model Name", placeholder="Enter Hugging Face model name (e.g., gpt2, distilgpt2)...")46    ],47    outputs="text",48    title="Prompt Refinement Tool",49    description="Input a prompt and model name to get a refined version that follows specific guidelines."50)51 52# Launch the app53if __name__ == "__main__":54    iface.launch()