Spaces:
Sleeping
Sleeping
File size: 4,707 Bytes
223c0a2 6900a40 223c0a2 31aa339 05199b4 223c0a2 31aa339 223c0a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import os
ADAPTER_PATH = "KirubaLS/Outreach_model_6"
BASE_MODEL_NAME = "google/gemma-2b"
MAX_SEQ_LENGTH = 1024
MAX_NEW_TOKENS = 250
def load_model_and_tokenizer(adapter_path, base_model_name, hf_token=None):
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
device = torch.device("cpu")
print(f"Using device: {device}")
try:
base_model = AutoModelForCausalLM.from_pretrained(
base_model_name,
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
token=hf_token,
device_map="auto"
)
print("Base model loaded successfully.")
except Exception as e:
print(f"Failed to load base model: {e}")
raise
try:
model = PeftModel.from_pretrained(
base_model,
adapter_path,
token=hf_token
)
print("PEFT adapter loaded and attached successfully.")
except Exception as e:
print(f"Failed to load PEFT adapter from {adapter_path}: {e}")
raise
try:
tokenizer = AutoTokenizer.from_pretrained(
base_model_name,
token=hf_token
)
print("Tokenizer loaded successfully.")
except Exception as e:
print(f"Failed to load tokenizer: {e}")
raise
model.eval()
return model, tokenizer, device
try:
model, tokenizer, device = load_model_and_tokenizer(ADAPTER_PATH, BASE_MODEL_NAME,hf_token=None)
except Exception as e:
print(f"Error loading model and tokenizer: {e}")
model = None
tokenizer = None
device = None
def generate_outreach_message(industry, company, prospect, designation, client_info, my_info):
if model is None or tokenizer is None:
return "Error: Model or tokenizer not loaded."
input_prompt = (
f"Industry: <industry>{industry}</industry>, "
f"Company: <company>{company}</company>, "
f"Prospect: <prospect>{prospect}</prospect>, "
f"Designation: <designation>{designation}</designation>, "
f"Client Info: <client_info>{client_info}</client_info>, "
f"My Info: <my_info>{my_info}</my_info>, "
f"Message Level: <message_level>First Level</message_level>: "
f"{industry}\nHi {prospect.split(' ')[0]},\n"
)
truncated_input_prompt = input_prompt
if len(tokenizer(input_prompt)['input_ids']) > MAX_SEQ_LENGTH:
truncated_input_prompt = tokenizer.decode(tokenizer(input_prompt)['input_ids'][:MAX_SEQ_LENGTH])
input_ids = tokenizer(
truncated_input_prompt,
return_tensors="pt",
padding=False,
truncation=True,
max_length=MAX_SEQ_LENGTH
).to(device)
with torch.no_grad():
output_tokens = model.generate(
input_ids["input_ids"],
max_new_tokens=MAX_NEW_TOKENS,
num_return_sequences=1,
no_repeat_ngram_size=2,
top_k=50,
top_p=0.95,
temperature=0.7,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
full_output_text = tokenizer.decode(output_tokens[0], skip_special_tokens=False)
input_prompt_in_output_index = full_output_text.find(truncated_input_prompt)
if input_prompt_in_output_index != -1:
generated_text = full_output_text[input_prompt_in_output_index + len(truncated_input_prompt):].strip()
else:
input_token_length = input_ids["input_ids"].shape[1]
generated_only_tokens = output_tokens[0, input_token_length:]
generated_text = tokenizer.decode(generated_only_tokens, skip_special_tokens=True).strip()
return generated_text
if __name__ == "__main__":
if model is not None and tokenizer is not None:
try:
df = pd.read_csv('dataset.csv')
data_row = df.loc[0]
generated_message = generate_outreach_message(
data_row['Industry'],
data_row['Company Name'],
data_row['Prospect Name'],
data_row['Designation'],
data_row['Client Info'],
data_row['My Info']
)
print("\n--- Test Generation ---")
print(generated_message)
print("-----------------------")
except FileNotFoundError:
print("dataset.csv not found. Cannot run test generation.")
else:
print("Model and tokenizer not loaded, cannot run test generation.") |