LeenAnabtawe commited on
Commit
02574a5
·
verified ·
1 Parent(s): 9a35ebc

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +41 -71
app.py CHANGED
@@ -1,95 +1,77 @@
1
  import gradio as gr
2
  from transformers import pipeline, AutoTokenizer
3
 
4
- # Load the StarChat beta model
5
- MODEL_NAME = "HuggingFaceH4/starchat-beta"
6
  tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
7
  generator = pipeline(
8
  "text-generation",
9
  model=MODEL_NAME,
10
  tokenizer=tokenizer,
11
- device="cpu", # Free tier uses CPU
12
  )
13
 
14
  def generate_code(natural_language_input):
15
  try:
16
- # Create a proper prompt for the model
17
- chat_prompt = [
18
- {"role": "system", "content": "You are a helpful AI assistant that generates Python code based on natural language descriptions."},
19
- {"role": "user", "content": f"Write Python code that: {natural_language_input}"}
20
- ]
21
 
22
- # Generate the prompt for the model
23
- prompt = tokenizer.apply_chat_template(chat_prompt, tokenize=False, add_generation_prompt=True)
24
-
25
- # Generate code with appropriate parameters for code generation
26
  generated = generator(
27
  prompt,
28
- max_new_tokens=256,
29
- temperature=0.2,
30
- top_p=0.95,
31
- do_sample=True,
32
  pad_token_id=tokenizer.eos_token_id,
33
- eos_token_id=tokenizer.eos_token_id,
34
  )
35
 
36
  # Extract and clean the generated code
37
  full_response = generated[0]['generated_text']
38
  code_response = full_response.replace(prompt, "").strip()
39
 
40
- # Sometimes the model adds non-code text - we'll try to extract just the code blocks
41
- if "```python" in code_response:
42
- code_response = code_response.split("```python")[1].split("```")[0]
43
- elif "```" in code_response:
44
- code_response = code_response.split("```")[1].split("```")[0]
45
 
46
  return code_response
47
  except Exception as e:
48
- return f"Error generating code: {str(e)}"
49
 
50
- # Gradio interface
51
- with gr.Blocks(theme=gr.themes.Soft()) as demo:
52
  gr.Markdown("""
53
- # 🚀 Text to Python Code Generator
54
- ### Powered by StarChat Beta
55
- Describe what you want the code to do in natural language, and get Python code!
56
  """)
57
 
58
  with gr.Row():
59
- with gr.Column(scale=3):
60
- input_text = gr.Textbox(
61
- label="Describe your code task",
62
- placeholder="e.g., 'create a Flask web server with two routes'",
63
- lines=3,
64
- max_lines=6
65
- )
66
- with gr.Row():
67
- generate_btn = gr.Button("Generate Code", variant="primary")
68
- clear_btn = gr.Button("Clear")
69
-
70
- gr.Markdown("**Tips:** Be specific in your description for better results.")
71
-
72
- with gr.Column(scale=4):
73
- output_code = gr.Code(
74
- label="Generated Python Code",
75
- language="python",
76
- interactive=True,
77
- lines=10
78
- )
79
- with gr.Row():
80
- copy_btn = gr.Button("Copy to Clipboard")
81
- refresh_btn = gr.Button("Try Again")
82
 
83
- # Examples
84
  examples = gr.Examples(
85
  examples=[
86
- ["create a function to calculate factorial recursively"],
87
- ["write code to download a file from URL and save it locally"],
88
- ["create a pandas DataFrame from a dictionary and plot a bar chart"],
89
- ["implement a simple REST API using FastAPI with GET and POST endpoints"]
90
  ],
91
  inputs=input_text,
92
- label="Click any example to try it"
93
  )
94
 
95
  # Button actions
@@ -104,18 +86,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
104
  inputs=None,
105
  outputs=[input_text, output_code]
106
  )
107
-
108
- refresh_btn.click(
109
- fn=generate_code,
110
- inputs=input_text,
111
- outputs=output_code
112
- )
113
-
114
- copy_btn.click(
115
- fn=lambda code: gr.Clipboard().copy(code),
116
- inputs=output_code,
117
- outputs=None,
118
- api_name="copy_code"
119
- )
120
 
121
- demo.launch(debug=True)
 
 
1
  import gradio as gr
2
  from transformers import pipeline, AutoTokenizer
3
 
4
+ # Load a smaller model that fits within free tier memory limits
5
+ MODEL_NAME = "Salesforce/codegen-350M-mono"
6
  tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
7
  generator = pipeline(
8
  "text-generation",
9
  model=MODEL_NAME,
10
  tokenizer=tokenizer,
11
+ device="cpu", # Force CPU usage for free tier
12
  )
13
 
14
  def generate_code(natural_language_input):
15
  try:
16
+ # Create a clear prompt for code generation
17
+ prompt = f"# Python code that {natural_language_input}\n#"
 
 
 
18
 
19
+ # Generate with conservative parameters to save memory
 
 
 
20
  generated = generator(
21
  prompt,
22
+ max_new_tokens=128, # Reduced from 256 to save memory
23
+ temperature=0.7,
24
+ top_p=0.9,
 
25
  pad_token_id=tokenizer.eos_token_id,
26
+ do_sample=True,
27
  )
28
 
29
  # Extract and clean the generated code
30
  full_response = generated[0]['generated_text']
31
  code_response = full_response.replace(prompt, "").strip()
32
 
33
+ # Remove any non-code text after generation
34
+ if "```" in code_response:
35
+ code_response = code_response.split("```")[0]
 
 
36
 
37
  return code_response
38
  except Exception as e:
39
+ return f"Error generating code: {str(e)}\n\n(Tip: Try a simpler description or break your request into smaller parts)"
40
 
41
+ # Streamlined Gradio interface to reduce memory footprint
42
+ with gr.Blocks(title="Text to Python Code Generator") as demo:
43
  gr.Markdown("""
44
+ # 🐍 Python Code Generator
45
+ *Free tier version using Salesforce/codegen-350M-mono*
 
46
  """)
47
 
48
  with gr.Row():
49
+ input_text = gr.Textbox(
50
+ label="Describe what you want the code to do",
51
+ placeholder="e.g., 'calculate fibonacci sequence up to N numbers'",
52
+ lines=3
53
+ )
54
+
55
+ with gr.Row():
56
+ generate_btn = gr.Button("Generate", variant="primary")
57
+ clear_btn = gr.Button("Clear")
58
+
59
+ output_code = gr.Code(
60
+ label="Generated Python Code",
61
+ language="python",
62
+ interactive=True
63
+ )
 
 
 
 
 
 
 
 
64
 
65
+ # Memory-friendly examples
66
  examples = gr.Examples(
67
  examples=[
68
+ ["print hello world"],
69
+ ["function to calculate square of a number"],
70
+ ["read a CSV file using pandas"],
71
+ ["simple Flask app with one route"]
72
  ],
73
  inputs=input_text,
74
+ label="Try these examples (keep requests simple)"
75
  )
76
 
77
  # Button actions
 
86
  inputs=None,
87
  outputs=[input_text, output_code]
88
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
89
 
90
+ # Launch with minimal resources
91
+ demo.launch(debug=False, show_error=True)