Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| MODEL_ID = "Jeppcode/ScalableLab2" | |
| SUBFOLDER = "merged-model-fp16" | |
| print(f"Loading model {MODEL_ID}/{SUBFOLDER} ...") | |
| # Tokenizer | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| MODEL_ID, subfolder=SUBFOLDER, | |
| ) | |
| # Model – fp16 and optimized for CPU | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID, | |
| subfolder=SUBFOLDER, | |
| dtype=torch.float16, | |
| low_cpu_mem_usage=True, | |
| device_map="cpu", | |
| ) | |
| model.eval() | |
| # Hardcoded system prompt | |
| SYSTEM_PROMPT = "You are a helpful, polite assistant. Give clear and structured explanations." | |
| def build_prompt(message, history): | |
| messages = [] | |
| messages.append({"role": "system", "content": SYSTEM_PROMPT}) | |
| for msg in history: | |
| role = msg.get("role") | |
| content = msg.get("content", "") | |
| if isinstance(content, list): | |
| texts = [] | |
| for block in content: | |
| if isinstance(block, dict) and block.get("type") == "text": | |
| texts.append(block.get("text", "")) | |
| else: | |
| texts.append(str(block)) | |
| text = "\n".join(t for t in texts if t) | |
| else: | |
| text = str(content) | |
| if text and role in ("user", "assistant", "system"): | |
| messages.append({"role": role, "content": text}) | |
| messages.append({"role": "user", "content": message}) | |
| prompt = tokenizer.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=True, | |
| ) | |
| return prompt | |
| def chat_fn(message, history, max_new_tokens): | |
| prompt = build_prompt(message, history) | |
| inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
| # Default values for removed sliders | |
| temperature = 0.7 | |
| top_p = 0.9 | |
| repetition_penalty = 1.1 | |
| gen_kwargs = { | |
| **inputs, | |
| "max_new_tokens": int(max_new_tokens), | |
| "pad_token_id": tokenizer.eos_token_id, | |
| "eos_token_id": tokenizer.eos_token_id, | |
| "repetition_penalty": float(repetition_penalty), | |
| "do_sample": True, | |
| "temperature": float(temperature), | |
| "top_p": float(top_p), | |
| } | |
| with torch.no_grad(): | |
| outputs = model.generate(**gen_kwargs) | |
| generated = tokenizer.decode( | |
| outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True, | |
| ).strip() | |
| return generated | |
| # --- Nature Theme Configuration --- | |
| nature_theme = gr.themes.Soft( | |
| primary_hue="green", | |
| secondary_hue="emerald", | |
| neutral_hue="stone", | |
| ).set( | |
| body_background_fill="#f5f7f5", | |
| block_background_fill="rgba(255, 255, 255, 0.9)", | |
| button_primary_background_fill="#2E7D32", | |
| border_color_primary="#4CAF50", | |
| ) | |
| custom_css = """ | |
| .gradio-container { | |
| background: url('https://images.unsplash.com/photo-1441974231531-c6227db76b6e?q=80&w=2560&auto=format&fit=crop') no-repeat center center fixed; | |
| background-size: cover; | |
| } | |
| footer {visibility: hidden} | |
| """ | |
| max_new_tokens_slider = gr.Slider( | |
| minimum=16, maximum=512, value=128, step=8, label="Response Length (Tokens)", | |
| ) | |
| # Instantiate ChatInterface WITHOUT theme/css arguments initially to avoid TypeError | |
| demo = gr.ChatInterface( | |
| fn=chat_fn, | |
| title="🌿 NatureChat Lab 2", | |
| description="Chat with the fine-tuned Llama model. Relax and enjoy the view.", | |
| additional_inputs=[max_new_tokens_slider], | |
| additional_inputs_accordion="Settings" | |
| ) | |
| # Apply theme and CSS manually after instantiation | |
| demo.theme = nature_theme | |
| demo.css = custom_css | |
| if __name__ == "__main__": | |
| demo.launch() |