Spaces:
Runtime error
Runtime error
File size: 1,513 Bytes
c82eb8a |
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 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import os
from dataclasses import dataclass, asdict
from ctransformers import AutoModelForCausalLM, AutoConfig
@dataclass
class GenerationConfig:
temperature: float
top_k: int
top_p: float
repetition_penalty: float
max_new_tokens: int
reset: bool
stream: bool
threads: int
stop: list[str]
def format_prompt(user_prompt: str):
return f"""### Instruction:
{user_prompt}
### Response:"""
def generate(llm: AutoModelForCausalLM,
generation_config: GenerationConfig,
prompt: str):
return llm(format_prompt(prompt), **asdict(generation_config))
def generate_code(prompt, config_path, model_name, max_tokens, temperature):
path = os.path.abspath(f"{config_path}/{model_name}.bin")
config = AutoConfig.from_pretrained(
os.path.abspath(config_path),
)
llm = AutoModelForCausalLM.from_pretrained(
path,
model_type="replit",
config=config,
)
generation_config = GenerationConfig(
temperature=temperature,
top_k=50,
top_p=0.9,
repetition_penalty=1.0,
max_new_tokens=max_tokens, # adjust as needed
reset=True, # reset history (cache)
stream=True, # streaming per word/token
threads=os.cpu_count(), # adjust for your CPU
stop=["<|endoftext|>"],
)
generator = generate(llm, generation_config, prompt)
output = ""
for word in generator:
print(word)
output += word
return output
|