| # demo.py - Quick demo of the model | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| model_name = "raihan-js/medllm-10m" | |
| print("Loading MedLLM...") | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| prompts = [ | |
| "Symptoms of diabetes include", | |
| "Treatment for high blood pressure", | |
| "The patient presents with" | |
| ] | |
| print("\nGenerating medical text:") | |
| for prompt in prompts: | |
| inputs = tokenizer(prompt, return_tensors="pt") | |
| outputs = model.generate( | |
| **inputs, | |
| max_length=50, | |
| do_sample=True, | |
| temperature=0.7, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| print(f"\nPrompt: {prompt}") | |
| print(f"Response: {response}") | |