File size: 5,247 Bytes
4f59c8a 21065ad 4f59c8a 21065ad 3f4d01e 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad 4f59c8a 21065ad | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from loguru import logger
# Configuration
# Using an open-access model that doesn't require authentication
# Options:
# 1. "HuggingFaceTB/SmolLM2-1.7B-Instruct" - Small, fast, no auth required
# 2. "meta-llama/Llama-3.2-3B-Instruct" - Better quality but requires HF login
# 3. "mistralai/Mistral-7B-Instruct-v0.3" - Larger, better quality
MODEL_REPO = "SaiBon99/llama-finetuned-merged"
logger.info(f"Loading model from {MODEL_REPO}...")
try:
# Load tokenizer
logger.info("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
logger.success(f"Tokenizer loaded successfully (vocab size: {tokenizer.vocab_size})")
# Load the model for CPU inference
logger.info("Loading model into memory (this may take a few minutes)...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_REPO,
dtype=torch.float32, # Use float32 for CPU, or bfloat16 to save memory
device_map="cpu",
low_cpu_mem_usage=True,
)
# Set to evaluation mode
model.eval()
param_count = sum(p.numel() for p in model.parameters())
logger.success(f"Model loaded successfully! Parameters: {param_count:,}")
except Exception as e:
logger.error(f"Error loading model: {e}")
logger.exception("Full traceback:")
model = None
tokenizer = None
def chat(message, history):
"""
Chat function that takes a message and chat history,
and returns the model's response.
"""
if model is None or tokenizer is None:
logger.error("No model or tokenizer loaded!")
return "Error: Model failed to load. Please check the logs."
logger.info(f"Received message: {message[:100]}..." if len(message) > 100 else f"Received message: {message}")
# Build the conversation using the chat template
messages = []
# Add history
if history:
logger.debug(f"Chat history has {len(history)} previous messages")
logger.info(f'history: {history}')
# Handle the new Gradio format where history is a list of message dictionaries
for msg in history:
role = msg.get('role', 'user')
# Extract text content from the message
content = msg.get('content', [])
if isinstance(content, list):
# Content is a list of content items, extract text from the first text item
content_text = ""
for item in content:
if isinstance(item, dict) and item.get('type') == 'text':
content_text = item.get('text', '')
break
else:
# Fallback if content is already a string
content_text = content
messages.append({"role": role, "content": content_text})
# Add current message
messages.append({"role": "user", "content": message})
# Generate response
try:
# Apply chat template and tokenize
logger.debug("Applying chat template...")
input_text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
inputs = tokenizer(input_text, return_tensors="pt")
input_length = inputs['input_ids'].shape[1]
logger.info(f"Input tokens: {input_length}")
# Generate
logger.info("Generating response...")
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=512,
temperature=0.7,
top_p=0.9,
repetition_penalty=1.1,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
)
# Decode only the new tokens
output_length = outputs.shape[1]
tokens_generated = output_length - input_length
logger.info(f"Generated {tokens_generated} new tokens")
response = tokenizer.decode(
outputs[0][input_length:],
skip_special_tokens=True
)
logger.success(f"Response generated successfully: {response[:100]}..." if len(response) > 100 else f"Response: {response}")
return response.strip()
except Exception as ex:
logger.error(f"Error generating response: {str(ex)}")
logger.exception("Full traceback:")
return f"Error generating response: {str(ex)}"
# Create Gradio ChatInterface
demo = gr.ChatInterface(
fn=chat,
title="Llama 3.2 3B Fine-tuned Chat",
description=f"Chat with Llama 3.2 3B using Hugging Face Transformers (CPU inference)\n\nModel: `{MODEL_REPO}`",
examples=[
"Explain what boolean operators are and how they work in programming.",
"What is the difference between short-circuit evaluation and normal evaluation?",
"Write a Python function to check if a number is prime.",
"Explain the concept of operator precedence with examples.",
],
)
if __name__ == "__main__":
logger.info("Starting Gradio interface...")
demo.launch(server_name="0.0.0.0") # Let Gradio auto-select available port
logger.info("Gradio interface launched successfully")
|