brunosan commited on
Commit
8a00478
1 Parent(s): 3381014

Upload gradio_interface.py

Browse files
Files changed (1) hide show
  1. gradio_interface.py +38 -0
gradio_interface.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #inference Gradio
2
+
3
+ import gradio as gr
4
+ import torch
5
+ from transformers import GPT2LMHeadModel, GPT2Tokenizer
6
+
7
+ # Load the fine-tuned model and tokenizer
8
+ model_path = 'brunosan/GPT2-impactscience'
9
+ tokenizer_path = 'brunosan/GPT2-impactscience'
10
+
11
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
12
+ tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_path)
13
+ model = GPT2LMHeadModel.from_pretrained(model_path).to(device)
14
+
15
+ # Define the generation function
16
+ def generate_text(prompt):
17
+ #remove trailing space if any
18
+ prompt = prompt.rstrip()
19
+ input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
20
+ attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device)
21
+ outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask,
22
+ max_length=100, num_beams=9,
23
+ no_repeat_ngram_size=2,
24
+ temperature=1.0, do_sample=True,
25
+ top_p=0.95, top_k=50)
26
+
27
+ generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
28
+ return generated_text
29
+
30
+ # Create a Gradio interface
31
+ input_text = gr.inputs.Textbox(lines=2, label="Enter the starting text")
32
+ output_text = gr.outputs.Textbox(label="Generated Text")
33
+
34
+ interface = gr.Interface(fn=generate_text, inputs=input_text, outputs=output_text,
35
+ title="GPT-2 Impact Science Text Generator", description="Generate text using a fine-tuned GPT-2 model onthe Impact Science book.")
36
+
37
+ if __name__ == "__main__":
38
+ interface.launch()