Views
No views yet
1from transformers import BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
2import torch
3from peft import PeftModel, PeftConfig
4from transformers import AutoModelForCausalLM, pipeline
5import logging
6# Suppress all warnings
7logging.getLogger("transformers").setLevel(logging.CRITICAL) #weird warning when using model for inference
8
9# Check if CUDA is available
10if torch.cuda.is_available():
11 num_devices = torch.cuda.device_count()
12 print(f"Number of available CUDA devices: {num_devices}")
13
14 for i in range(num_devices):
15 device_name = torch.cuda.get_device_name(i)
16 print(f"\nDevice {i}: {device_name}")
17else:
18 print("CUDA is not available.")
19# Specify the device (0 for GPU or -1 for CPU)
20device = 0 if torch.cuda.is_available() else -1
21
22config = PeftConfig.from_pretrained("smartinez1/Llama-3.1-8B-FINLLM")
23base_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
24model = PeftModel.from_pretrained(base_model, "smartinez1/Llama-3.1-8B-FINLLM")
25# Load the tokenizer associated with the base model
26tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")
27# Define the unique padding token for fine-tuning
28custom_pad_token = "<|finetune_right_pad_id|>"
29tokenizer.add_special_tokens({'pad_token': custom_pad_token})
30pad_token_id = tokenizer.pad_token_id
31
32# Set up the text generation pipeline with the PEFT model, specifying the device
33generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
34
35# List of user inputs
36user_inputs = [
37 "Provide a link for Regulation A (Extensions of Credit by Federal Reserve Banks) law",
38 "Define the following term: Insurance Scores.",
39 "Expand the following acronym into its full form: ESCB.",
40 "Provide a concise answer to the following question: Which countries currently have bilateral FTAs in effect with the U.S.?",
41 """Given the following text, only list the following for each: specific Organizations, Legislations, Dates, Monetary Values,
42 and Statistics When can counterparties start notifying the national competent authorities (NCAs) of their intention to apply
43 the reporting exemption in accordance with Article 9(1) EMIR, as amended by Regulation 2019/834?""",
44 "Provide a concise answer to the following question: What type of license is the Apache License, Version 2.0?"
45]
46
47# Define the prompt template
48prompt_template = """Below is an instruction that describes a task. Write a response that appropriately completes the request.
49
50### Instruction:
51{}
52
53### Answer:
54"""
55
56# Loop over each user input and generate a response
57for user_input in user_inputs:
58 # Format the user input into the prompt
59 prompt = prompt_template.format(user_input)
60
61 # Generate a response from the model
62 response = generator(prompt, max_length=200, num_return_sequences=1, do_sample=True)
63
64 # Extract and clean up the AI's response
65 response_str = response[0]['generated_text'].split('### Answer:')[1].strip()
66 cut_ind = response_str.find("#") # Remove extra information after the response
67 response_str = response_str[:cut_ind].strip() if cut_ind != -1 else response_str
68
69 # Display the AI's response
70 print(f"User: {user_input}")
71 print(f"AI: {response_str}")
72 print("-" * 50) # Separator for clarity