Main.py

# 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)