Spaces:
Build error
Build error
File size: 2,835 Bytes
7f5f0ca | 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 | import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig
import torch
import logging
import gc
def clear_memory():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
def load_model():
# Force CPU for stability
device = torch.device("cpu")
# Load base model
base_model = AutoModelForCausalLM.from_pretrained(
"microsoft/phi-2",
torch_dtype=torch.float32,
trust_remote_code=True,
device_map=None,
low_cpu_mem_usage=True
)
# Load LoRA configuration
peft_config = PeftConfig.from_pretrained("phi2-oasst-qlora-final")
# Load and merge LoRA weights
model = PeftModel.from_pretrained(
base_model,
"phi2-oasst-qlora-final",
device_map=None,
torch_dtype=torch.float32
)
# Merge LoRA weights with base model
model = model.merge_and_unload()
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
"microsoft/phi-2",
trust_remote_code=True
)
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
def generate_response(prompt):
# Format prompt
full_prompt = f"Instruct: {prompt}\nOutput:"
# Tokenize
inputs = tokenizer(
full_prompt,
return_tensors="pt",
padding=False,
truncation=True,
max_length=256
)
# Generate
with torch.inference_mode():
outputs = model.generate(
**inputs,
max_length=128,
min_length=10,
temperature=0.7,
top_p=0.9,
num_return_sequences=1,
pad_token_id=tokenizer.eos_token_id,
do_sample=True,
no_repeat_ngram_size=2,
use_cache=True
)
# Decode response
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Remove the prompt from response
response = response.replace(full_prompt, "").strip()
return response
# Load model and tokenizer globally
print("Loading model...")
model, tokenizer = load_model()
print("Model loaded!")
# Create Gradio interface
iface = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(
lines=3,
placeholder="Enter your instruction here...",
label="Input"
),
outputs=gr.Textbox(
lines=5,
label="Generated Response"
),
title="Phi-2 Fine-tuned Assistant",
description="A fine-tuned version of Phi-2 on the OpenAssistant dataset",
examples=[
["Explain what is Python in one sentence."],
["How do neural networks learn through backpropagation?"],
["Write a short poem about coding."]
]
)
if __name__ == "__main__":
iface.launch() |