Testing / Main.py
mohommed's picture
Create Main.py
930e89e verified
Raw History Blame Contribute Delete
1.31 kB
# Install dependencies
!pip install transformers accelerate peft
# Import libraries
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel, PeftConfig
import torch
# Load base model and PEFT adapter
base_model_name = "EleutherAI/gpt-neo-1.3B"
peft_model_name = "mohommed/Testing"
# Load tokenizer and base model
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
base_model = AutoModelForCausalLM.from_pretrained(base_model_name, device_map="auto", torch_dtype=torch.float16)
# Apply the PEFT adapter
model = PeftModel.from_pretrained(base_model, peft_model_name)
model.eval()
# Define generation function
def generate_article(prompt, max_new_tokens=700):
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=0.8,
top_p=0.95,
do_sample=True,
eos_token_id=tokenizer.eos_token_id,
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Example prompt
prompt = "Write a 600-word Medium article about how to stay motivated as a writer. Use a personal tone and clear structure."
# Generate article
output = generate_article(prompt)
print(output)