Spaces:
Paused
Paused
| import gradio as gr | |
| import transformers as t | |
| import torch | |
| import peft | |
| # Load your fine-tuned model and tokenizer | |
| tokenizer = t.AutoTokenizer.from_pretrained("NousResearch/Llama-2-7b-hf") | |
| model = t.AutoModelForCausalLM.from_pretrained("NousResearch/Llama-2-7b-hf") | |
| tokenizer.pad_token_id = 0 | |
| config = peft.LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"], lora_dropout=0.005, bias="none", task_type="CAUSAL_LM") | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = peft.get_peft_model(model, config).to(device) | |
| peft.set_peft_model_state_dict(model, torch.load(f".weights/adapter_model.bin")) | |
| # Define a prediction function | |
| def generate_article(title): | |
| prompt = f"Below is a title for an article. Write an article that appropriately suits the title: \n\n### Title:\n{title}\n\n### Article:\n" | |
| pipe = t.pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=1000) | |
| output = pipe([prompt]) | |
| generated_article = output[0][0]["generated_text"] | |
| return generated_article | |
| # Create a Gradio interface | |
| iface = gr.Interface( | |
| fn=generate_article, | |
| inputs=gr.inputs.Textbox(lines=2, placeholder="Enter Article Title Here"), | |
| outputs="text", | |
| title="Article Generator", | |
| description="Enter a title to generate an article." | |
| ) | |
| # Launch the app | |
| iface.launch() | |