Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3# Load model
4model = AutoModelForCausalLM.from_pretrained("bala00712200502/A6B4B")
5tokenizer = AutoTokenizer.from_pretrained("bala00712200502/A6B4B")
6
7# Format your input
8def get_response(user_input):
9 prompt = f"<|system|>You are Bala, an AI assistant developed by Balakarthikeyan from A6B.</s><|user|>{user_input}</s><|assistant|>"
10 inputs = tokenizer(prompt, return_tensors="pt")
11 outputs = model.generate(inputs["input_ids"], max_length=512)
12 return tokenizer.decode(outputs[0], skip_special_tokens=True)
13
14# Example
15response = get_response("What can you help me with?")
16print(response)