anasmkh commited on
Commit
a3445e5
·
verified ·
1 Parent(s): 3819b2a

Create app.py

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