Spaces:
Sleeping
Sleeping
| 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.") |