Add Be.FM-8B chat interface with PEFT adapter

app.py

@@ -0,0 +1,105 @@
+# app.py
+import os
+import torch
+import spaces
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
+
+BASE_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+PEFT_MODEL_ID = "befm/Be.FM-8B"
+
+USE_PEFT = True
+try:
+    from peft import PeftModel, PeftConfig  # noqa
+except Exception:
+    USE_PEFT = False
+    print("[WARN] 'peft' not installed; running base model only.")
+
+def load_model_and_tokenizer():
+    if HF_TOKEN is None:
+        raise RuntimeError(
+            "HF_TOKEN is not set. Add it in Space → Settings → Secrets. "
+            "Also ensure your account has access to the gated base model."
+        )
+    dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+    tok = AutoTokenizer.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    base = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL_ID,
+        device_map="auto" if torch.cuda.is_available() else None,
+        torch_dtype=dtype,
+        token=HF_TOKEN,
+    )
+
+    if USE_PEFT:
+        try:
+            _ = PeftConfig.from_pretrained(PEFT_MODEL_ID, use_auth_token=HF_TOKEN)
+            model = PeftModel.from_pretrained(base, PEFT_MODEL_ID, use_auth_token=HF_TOKEN)
+            print(f"[INFO] Loaded PEFT adapter: {PEFT_MODEL_ID}")
+            return model, tok
+        except Exception as e:
+            print(f"[WARN] Failed to load PEFT adapter: {e}")
+            return base, tok
+    return base, tok
+
+model, tokenizer = load_model_and_tokenizer()
+DEVICE = model.device
+
+@spaces.GPU
+@torch.inference_mode()
+def generate_response(prompt: str, max_new_tokens=512, temperature=0.7, top_p=0.9) -> str:
+    enc = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    enc = {k: v.to(DEVICE) for k, v in enc.items()}
+    out = model.generate(
+        **enc,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    return tokenizer.decode(out[0], skip_special_tokens=True)
+
+def chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p):
+    # Build a simple conversation string
+    conv = []
+    if system_prompt:
+        conv.append(f"system: {system_prompt}")
+    for u, a in (history or []):
+        if u:
+            conv.append(f"user: {u}")
+        if a:
+            conv.append(f"assistant: {a}")
+    if message:
+        conv.append(f"user: {message}")
+    prompt = "\n".join(conv) + "\nassistant:"
+    reply = generate_response(
+        prompt,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+    )
+    # Strip trailing
+    if "assistant:" in reply:
+        reply = reply.split("assistant:")[-1].strip()
+    return reply
+
+demo = gr.ChatInterface(
+    fn=lambda message, history, system_prompt, max_new_tokens, temperature, top_p:
+        chat_fn(message, history, system_prompt, max_new_tokens, temperature, top_p),
+    additional_inputs=[
+        gr.Textbox(label="System prompt (optional)", placeholder="You are Be.FM assistant...", lines=2),
+        gr.Slider(16, 2048, value=512, step=16, label="max_new_tokens"),
+        gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="temperature"),
+        gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p"),
+    ],
+    title="Be.FM-8B (PEFT) on Meta-Llama-3.1-8B-Instruct",
+    description="Chat interface using Meta-Llama-3.1-8B-Instruct with PEFT adapter befm/Be.FM-8B."
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.0.0
+transformers>=4.30.0
+peft>=0.4.0
+spaces
+accelerate