| # inference.py | |
| from transformers import pipeline | |
| # Initialize the pipeline | |
| model_name = "Mr-Vicky-01/Gemma-2B-Finetuined-pythonCode" | |
| pipe = pipeline("text2text-generation", model=model_name) | |
| # Function to generate code based on user input | |
| def generate_code(prompt): | |
| # Use the pipeline to generate text | |
| generated_code = pipe(prompt) | |
| return generated_code[0]['generated_text'] # Extract the generated text | |
| if __name__ == "__main__": | |
| # Example prompt for code generation | |
| user_prompt = "Write a Python function to calculate the Fibonacci sequence." | |
| result = generate_code(user_prompt) | |
| print("Generated Code:\n", result) | |