Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
3from peft import PeftModel
4
5base_model_id = "Adit1Sharma/fineweb-1.5b-checkpoints"
6adapter_id = "Adit1Sharma/fineweb-1.5b-final-sft-adapter"
7
8# 1. Load the tokenizer
9tokenizer = AutoTokenizer.from_pretrained("gpt2")
10
11# 2. Load the base model
12print("Loading base model...")
13base_model = AutoModelForCausalLM.from_pretrained(
14 base_model_id,
15 trust_remote_code=True,
16 device_map="auto"
17)
18
19# 3. Apply the SFT Adapter dynamically
20print("Applying SFT adapter...")
21model = PeftModel.from_pretrained(
22 base_model,
23 adapter_id
24)
25
26# 4. Set up the pipeline with sampling parameters
27model.generation_config.max_length = None # kill the stale max_length=20
28
29pipe = pipeline(
30 task="text-generation",
31 model=model,
32 tokenizer=tokenizer,
33 max_new_tokens=200,
34 pad_token_id=tokenizer.eos_token_id,
35 do_sample=True,
36 temperature=0.6,
37 top_p=0.6,
38 top_k=40,
39 repetition_penalty=1.1,
40 no_repeat_ngram_size=3,
41)
42# 5. Generate Text
43prompt = "Artificial intelligence is"
44output = pipe(prompt)
45print(output[0]['generated_text'])
46