engineer2411/developer

1from transformers import AutoTokenizer, AutoModelForCausalLM2import os3 4hf_token = os.getenv("HF_KEY")5model_id = "engineer2411/Jarvis_TinyLlama-1.1B-Chat_v1"6 7def get_inference(text, model, tokenizer, max_input_tokens=50, max_output_tokens=50):8	# Tokenize9	input_ids = tokenizer.encode(10		text,11		return_tensors="pt",12		truncation=True,13		max_length=max_input_tokens14	)15 16	# Generate17	device = model.device18	generated_tokens_with_prompt = model.generate(19		input_ids=input_ids.to(device),20		max_length=max_output_tokens21	)22 23	# Decode24	generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)25 26	# Strip the prompt27	generated_text_answer = generated_text_with_prompt[0][len(text):]28	return generated_text_answer29 30def get_jarvis_response(question):31	try:32		tokenizer = AutoTokenizer.from_pretrained(model_id)33		tokenizer.pad_token = tokenizer.eos_token34		model = AutoModelForCausalLM.from_pretrained(model_id)35		return get_inference(question, model, tokenizer)36	except Exception as e:37		print(f"Error during inference: {e}")38		return e