gemma-2b-telecom-sft / inference.py
ligaments-dev's picture
Add inference script for deployed model
1c38647 verified
from transformers import pipeline
# After training, the model is available at:
MODEL_ID = "ligaments-dev/gemma-2b-telecom-sft"
pipe = pipeline("text-generation", model=MODEL_ID, device_map="auto")
messages = [
{"role": "user", "content": "My internet has been down for two hours. Can you help me troubleshoot?"}
]
response = pipe(messages, max_new_tokens=256, do_sample=True, temperature=0.7)
print(response[0]["generated_text"])