Canstralian commited on
Commit
19fb0b2
·
verified ·
1 Parent(s): 1892c23

Create inference.py

Browse files
Files changed (1) hide show
  1. inference.py +19 -0
inference.py ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # inference.py
2
+
3
+ from transformers import pipeline
4
+
5
+ # Initialize the pipeline
6
+ model_name = "Mr-Vicky-01/Gemma-2B-Finetuined-pythonCode"
7
+ pipe = pipeline("text2text-generation", model=model_name)
8
+
9
+ # Function to generate code based on user input
10
+ def generate_code(prompt):
11
+ # Use the pipeline to generate text
12
+ generated_code = pipe(prompt)
13
+ return generated_code[0]['generated_text'] # Extract the generated text
14
+
15
+ if __name__ == "__main__":
16
+ # Example prompt for code generation
17
+ user_prompt = "Write a Python function to calculate the Fibonacci sequence."
18
+ result = generate_code(user_prompt)
19
+ print("Generated Code:\n", result)