| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
| from peft import PeftModel |
|
|
| BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B" |
| ADAPTER = "Nayan7r839/EchooaiForge" |
|
|
| print("Loading tokenizer...") |
| tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL) |
|
|
| print("Loading base model...") |
| base_model = AutoModelForCausalLM.from_pretrained( |
| BASE_MODEL, |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, |
| device_map="auto" |
| ) |
|
|
| print("Loading EchooaiForge adapter...") |
| model = PeftModel.from_pretrained( |
| base_model, |
| ADAPTER |
| ) |
|
|
| model.eval() |
|
|
| print("EchooaiForge loaded!") |
|
|
|
|
| def generate_code(prompt: str): |
| text = f"""### Instruction: |
| {prompt} |
| |
| ### Response: |
| """ |
|
|
| inputs = tokenizer( |
| text, |
| return_tensors="pt" |
| ).to(model.device) |
|
|
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=512, |
| temperature=0.2, |
| do_sample=True, |
| top_p=0.9, |
| repetition_penalty=1.05 |
| ) |
|
|
| result = tokenizer.decode( |
| outputs[0], |
| skip_special_tokens=True |
| ) |
|
|
| if "### Response:" in result: |
| result = result.split("### Response:", 1)[1] |
|
|
| return result.strip() |
