File size: 649 Bytes
19fb0b2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
# 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)
|