Spaces:
Sleeping
Sleeping
File size: 1,453 Bytes
d9ded01 b2c1928 ea0af80 b2c1928 d9ded01 8374669 d36dc81 ba41c7f d9ded01 ba41c7f 948bd8f ba41c7f d36dc81 948bd8f eb66cb5 e2116c0 ea0af80 948bd8f e2116c0 ad9f174 e2116c0 ad9f174 e2116c0 948bd8f ad9f174 e2116c0 948bd8f ad9f174 e2116c0 948bd8f e2116c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
MODEL_NAME = "bigcode/starcoderbase-1b"
HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
device = "cpu"
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
token=HF_TOKEN,
torch_dtype=torch.float32,
trust_remote_code=True
).to(device)
def generate_code(prompt: str, max_tokens: int = 256):
# 🛠 **Improve Speed & Support Multi-language**
formatted_prompt = f"{prompt}\n### Code:\n" # Hint that code follows
inputs = tokenizer(
formatted_prompt,
return_tensors="pt",
truncation=True,
max_length=512 # ⏩ Reduce max length to speed up processing
).to(device)
output = model.generate(
**inputs,
max_new_tokens=max_tokens,
pad_token_id=tokenizer.pad_token_id,
do_sample=True,
top_p=0.90, # 🎯 Prioritize better predictions
temperature=0.6 # 🔥 More deterministic output
)
generated_code = tokenizer.decode(output[0], skip_special_tokens=True)
# Remove the input prompt from the output
if generated_code.startswith(formatted_prompt):
generated_code = generated_code[len(formatted_prompt):]
return generated_code.strip()
|