Harsh_GRPO / app.py
TheDeepDas's picture
Upload folder using huggingface_hub
0af2a1b verified
Raw History Blame Contribute Delete
2.85 kB
from pathlib import Path
import gradio as gr
import spaces
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
BASE_MODEL = "Qwen/Qwen2.5-Coder-3B-Instruct"
ADAPTER_DIR = Path(__file__).resolve().parent
def load_model():
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_DIR)
use_cuda = torch.cuda.is_available()
dtype = torch.float16 if use_cuda else torch.float32
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
torch_dtype=dtype,
low_cpu_mem_usage=True,
)
if use_cuda:
# ZeroGPU requires the model to be placed on CUDA during app startup.
base_model = base_model.to("cuda")
else:
base_model = base_model.to("cpu")
model = PeftModel.from_pretrained(base_model, ADAPTER_DIR)
model.eval()
return model, tokenizer
model, tokenizer = load_model()
def gpu_duration(prompt: str, max_new_tokens: int, temperature: float) -> int:
# Allow time for long requests while keeping typical calls in the fast queue.
return min(120, max(30, 15 + int(max_new_tokens) // 4))
@spaces.GPU(duration=gpu_duration)
@torch.inference_mode()
def predict(prompt: str, max_new_tokens: int, temperature: float) -> str:
if not prompt.strip():
raise gr.Error("Enter a prompt before requesting an inference.")
inputs = tokenizer.apply_chat_template(
[{"role": "user", "content": prompt}],
tokenize=True,
add_generation_prompt=True,
return_tensors="pt",
return_dict=True,
)
inputs = inputs.to(model.get_input_embeddings().weight.device)
generation_options = {
"max_new_tokens": int(max_new_tokens),
"eos_token_id": tokenizer.eos_token_id,
"pad_token_id": tokenizer.pad_token_id or tokenizer.eos_token_id,
"do_sample": temperature > 0,
}
if temperature > 0:
generation_options["temperature"] = temperature
output = model.generate(**inputs, **generation_options)
generated_tokens = output[0, inputs["input_ids"].shape[1] :]
return tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
demo = gr.Interface(
fn=predict,
inputs=[
gr.Textbox(label="Prompt", lines=5, placeholder="Ask for help writing or explaining code..."),
gr.Slider(minimum=16, maximum=512, value=128, step=16, label="Maximum new tokens"),
gr.Slider(minimum=0, maximum=1.5, value=0.7, step=0.1, label="Temperature (0 = deterministic)"),
],
outputs=gr.Textbox(label="Model response", lines=10),
title="GRPO Qwen2.5 Coder",
description="Qwen2.5-Coder-3B-Instruct fine-tuned with a GRPO LoRA adapter.",
api_name="predict",
)
if __name__ == "__main__":
demo.launch()