Spaces:
Sleeping
Sleeping
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import torch | |
# Use a basic small code model — change to any Hugging Face model you want | |
MODEL_NAME = "Salesforce/codegen-350M-multi" | |
def load_model(): | |
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME) | |
return tokenizer, model | |
def generate_code(tokenizer, model, prompt, max_length=128): | |
inputs = tokenizer(prompt, return_tensors="pt") | |
outputs = model.generate(**inputs, max_length=max_length) | |
return tokenizer.decode(outputs[0], skip_special_tokens=True) | |