Spaces:
Running
Running
File size: 4,042 Bytes
218e523 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
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) |