Spaces:
Running
Running
| import gradio as gr | |
| from transformers import pipeline, AutoTokenizer | |
| # Load the StarChat beta model | |
| MODEL_NAME = "HuggingFaceH4/starchat-beta" | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| generator = pipeline( | |
| "text-generation", | |
| model=MODEL_NAME, | |
| tokenizer=tokenizer, | |
| device="cpu", # Free tier uses CPU | |
| ) | |
| def generate_code(natural_language_input): | |
| try: | |
| # Create a proper prompt for the model | |
| chat_prompt = [ | |
| {"role": "system", "content": "You are a helpful AI assistant that generates Python code based on natural language descriptions."}, | |
| {"role": "user", "content": f"Write Python code that: {natural_language_input}"} | |
| ] | |
| # Generate the prompt for the model | |
| prompt = tokenizer.apply_chat_template(chat_prompt, tokenize=False, add_generation_prompt=True) | |
| # Generate code with appropriate parameters for code generation | |
| generated = generator( | |
| prompt, | |
| max_new_tokens=256, | |
| temperature=0.2, | |
| top_p=0.95, | |
| do_sample=True, | |
| pad_token_id=tokenizer.eos_token_id, | |
| eos_token_id=tokenizer.eos_token_id, | |
| ) | |
| # Extract and clean the generated code | |
| full_response = generated[0]['generated_text'] | |
| code_response = full_response.replace(prompt, "").strip() | |
| # Sometimes the model adds non-code text - we'll try to extract just the code blocks | |
| if "```python" in code_response: | |
| code_response = code_response.split("```python")[1].split("```")[0] | |
| elif "```" in code_response: | |
| code_response = code_response.split("```")[1].split("```")[0] | |
| return code_response | |
| except Exception as e: | |
| return f"Error generating code: {str(e)}" | |
| # Gradio interface | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 🚀 Text to Python Code Generator | |
| ### Powered by StarChat Beta | |
| Describe what you want the code to do in natural language, and get Python code! | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| input_text = gr.Textbox( | |
| label="Describe your code task", | |
| placeholder="e.g., 'create a Flask web server with two routes'", | |
| lines=3, | |
| max_lines=6 | |
| ) | |
| with gr.Row(): | |
| generate_btn = gr.Button("Generate Code", variant="primary") | |
| clear_btn = gr.Button("Clear") | |
| gr.Markdown("**Tips:** Be specific in your description for better results.") | |
| with gr.Column(scale=4): | |
| output_code = gr.Code( | |
| label="Generated Python Code", | |
| language="python", | |
| interactive=True, | |
| lines=10 | |
| ) | |
| with gr.Row(): | |
| copy_btn = gr.Button("Copy to Clipboard") | |
| refresh_btn = gr.Button("Try Again") | |
| # Examples | |
| examples = gr.Examples( | |
| examples=[ | |
| ["create a function to calculate factorial recursively"], | |
| ["write code to download a file from URL and save it locally"], | |
| ["create a pandas DataFrame from a dictionary and plot a bar chart"], | |
| ["implement a simple REST API using FastAPI with GET and POST endpoints"] | |
| ], | |
| inputs=input_text, | |
| label="Click any example to try it" | |
| ) | |
| # Button actions | |
| generate_btn.click( | |
| fn=generate_code, | |
| inputs=input_text, | |
| outputs=output_code | |
| ) | |
| clear_btn.click( | |
| fn=lambda: ("", ""), | |
| inputs=None, | |
| outputs=[input_text, output_code] | |
| ) | |
| refresh_btn.click( | |
| fn=generate_code, | |
| inputs=input_text, | |
| outputs=output_code | |
| ) | |
| copy_btn.click( | |
| fn=lambda code: gr.Clipboard().copy(code), | |
| inputs=output_code, | |
| outputs=None, | |
| api_name="copy_code" | |
| ) | |
| demo.launch(debug=True) |