File size: 816 Bytes
a3445e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
model_name = "anasmkh/customized_llama3.1_8b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
generator = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=64,
temperature=1.5,
min_p=0.1
)
def generate_response(prompt):
messages = [
{"role": "user", "content": prompt},
]
response = generator(messages)[0]['generated_text']
return response.split("<|end_header_id|>")[1].strip()
demo = gr.Interface(
fn=generate_response,
inputs=gr.Textbox(lines=5, label="Enter your prompt"),
outputs=gr.Textbox(label="Model Response")
)
demo.launch()
|