Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4model = AutoModelForCausalLM.from_pretrained("Pravincoder/Pythia-legal-finetuned-llm")
5tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m")
6
7def inference(text, model, tokenizer, max_input_tokens=1000, max_output_tokens=200):
8 input_ids = tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=max_input_tokens)
9 device = model.device
10 generated_tokens_with_prompt = model.generate(input_ids=input_ids.to(device), max_length=max_output_tokens)
11 generated_text_with_prompt = tokenizer.batch_decode(generated_tokens_with_prompt, skip_special_tokens=True)
12 generated_text_answer = generated_text_with_prompt[0][len(text):]
13 return generated_text_answer
14
15system_message = "Welcome to the medical AI assistant."
16user_message = "What are the symptoms of influenza?"
17generated_response = inference(system_message, user_message, model, tokenizer)
18print("Generated Response:", generated_response)