Views
No views yet
meta-llama/Llama-3.2-1B-Instruct1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from peft import PeftModel
4
5
6def load_model_and_tokenizer(base_model_name: str, lora_checkpoint: str):
7 """
8 Load the tokenizer and fine-tuned LoRA model.
9
10 Args:
11 base_model_name (str): Name of the base model from Hugging Face.
12 lora_checkpoint (str): Path or Hugging Face repository of the LoRA fine-tuned model.
13
14 Returns:
15 model: The LoRA fine-tuned model.
16 tokenizer: Tokenizer for the model.
17 """
18 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
19 base_model = AutoModelForCausalLM.from_pretrained(
20 base_model_name,
21 device_map="auto",
22 torch_dtype=torch.float16,
23 )
24 model = PeftModel.from_pretrained(base_model, lora_checkpoint)
25
26 # Set the padding token to the EOS token if not defined
27 if tokenizer.pad_token is None:
28 tokenizer.pad_token = tokenizer.eos_token
29 model.config.pad_token_id = tokenizer.pad_token_id
30
31 model.eval()
32 return model, tokenizer
33
34
35def generate_response(
36 model,
37 tokenizer,
38 input_text: str,
39 prompt_template: str = None,
40 max_length: int = 512,
41 num_beams: int = 4,
42 temperature: float = 0.7,
43 top_k: int = 50,
44 top_p: float = 0.9,
45 repetition_penalty: float = 1.2,
46 do_sample: bool = True,
47):
48 """
49 Generate a response from the model given an input prompt.
50
51 Args:
52 model: The LoRA fine-tuned model.
53 tokenizer: Tokenizer for the model.
54 input_text (str): User input or prompt.
55 prompt_template (str): Template for generating responses (optional).
56 max_length (int): Maximum length of the response.
57 num_beams (int): Number of beams for beam search.
58 temperature (float): Sampling temperature.
59 top_k (int): Top-k sampling parameter.
60 top_p (float): Top-p sampling parameter.
61 repetition_penalty (float): Penalty for word repetition.
62 do_sample (bool): Whether to enable sampling.
63
64 Returns:
65 str: Generated response from the model.
66 """
67 if prompt_template:
68 input_text = prompt_template.format(input_text=input_text)
69
70 inputs = tokenizer(
71 input_text,
72 return_tensors="pt",
73 padding=True,
74 truncation=True,
75 )
76
77 with torch.no_grad():
78 output = model.generate(
79 input_ids=inputs["input_ids"].to(model.device),
80 attention_mask=inputs["attention_mask"].to(model.device),
81 max_length=max_length,
82 num_beams=num_beams,
83 temperature=temperature if do_sample else None,
84 top_k=top_k if do_sample else None,
85 top_p=top_p if do_sample else None,
86 repetition_penalty=repetition_penalty,
87 do_sample=do_sample,
88 early_stopping=True,
89 pad_token_id=tokenizer.pad_token_id,
90 eos_token_id=tokenizer.eos_token_id,
91 )
92
93 response = tokenizer.decode(output[0], skip_special_tokens=True)
94 return response.strip()
95
96
97def main():
98 """
99 Main function to load the model, provide user prompts, and generate responses.
100 """
101 # Configuration
102 BASE_MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
103 LORA_CHECKPOINT = "AIAlbus/EffiLLaMA" # Replace with your Hugging Face model repo or local path
104 PROMPT_TEMPLATE = """Analyze the given question based on facts established in the Harry Potter series canon.
105
106Rules:
1071. Use only information from the books, films, or official sources like interviews with J.K. Rowling.
1082. Avoid inventing details, characters, or events not present in canon.
1093. If analysis or interpretation is provided, explicitly state it as such.
110
111Question: {input_text}
112
113Factual analysis:
114"""
115
116 print("Loading model and tokenizer...")
117 model, tokenizer = load_model_and_tokenizer(BASE_MODEL_NAME, LORA_CHECKPOINT)
118 print("Model and tokenizer loaded successfully!")
119
120 print("\n--- Welcome to EffiLLaMA Inference Script ---")
121 print("Enter your prompt below (type 'exit' to quit):\n")
122
123 while True:
124 user_input = input("Your Prompt: ").strip()
125 if user_input.lower() == "exit":
126 print("Exiting the inference script. Goodbye!")
127 break
128
129 print("\nGenerating response...\n")
130 response = generate_response(
131 model=model,
132 tokenizer=tokenizer,
133 input_text=user_input,
134 prompt_template=PROMPT_TEMPLATE,
135 )
136 print(f"Response:\n{response}")
137 print("-" * 80)
138
139
140if __name__ == "__main__":
141 main()