Wtzwho commited on
Commit
a79e48c
1 Parent(s): 663b5a6

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, pipeline
3
+ import torch
4
+
5
+ # Initialize the model and tokenizer
6
+ model_name = "Wtzwho/Prometh-MOEM-V.01"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ text_generation_pipeline = pipeline(
9
+ "text-generation",
10
+ model=model_name,
11
+ model_kwargs={"torch_dtype": torch.float16, "load_in_4bit": True},
12
+ )
13
+
14
+ def generate_text(user_input):
15
+ messages = [{"role": "user", "content": user_input}]
16
+ prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
17
+ outputs = text_generation_pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
18
+ return outputs[0]["generated_text"]
19
+
20
+ # Create the Gradio interface
21
+ iface = gr.Interface(
22
+ fn=generate_text,
23
+ inputs=gr.inputs.Textbox(lines=2, placeholder="Type your question here..."),
24
+ outputs=gr.outputs.Textbox(),
25
+ title="Prometh-MOEM Text Generation",
26
+ description="A text generation model that understands your queries and generates concise, informative responses."
27
+ )
28
+
29
+ # Run the interface
30
+ iface.launch()