RayBe commited on
Commit
8920931
·
verified ·
1 Parent(s): e276d9f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+ import torch
3
+ import gradio as gr
4
+ from transformers import T5Tokenizer, T5ForConditionalGeneration
5
+
6
+ # Load the fine-tuned model
7
+ model_name = "./t5-finetuned-final"
8
+ tokenizer = T5Tokenizer.from_pretrained(model_name)
9
+ model = T5ForConditionalGeneration.from_pretrained(model_name)
10
+
11
+ # Move model to GPU if available
12
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
13
+ model.to(device)
14
+
15
+ # Define the function for inference
16
+ def generate_command(input_command):
17
+ prompt = "extract: " + input_command
18
+ input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
19
+ output_ids = model.generate(input_ids, max_length=128, num_beams=5, early_stopping=True)
20
+ return tokenizer.decode(output_ids[0], skip_special_tokens=True)
21
+
22
+ # Create a Gradio interface
23
+ iface = gr.Interface(
24
+ fn=generate_command,
25
+ inputs=gr.Textbox(lines=2, placeholder="Enter a command..."),
26
+ outputs=gr.Textbox(label="Extracted JSON Output"),
27
+ title="T5 Fine-Tuned Command Extractor",
28
+ description="Enter a command, and the fine-tuned T5 model will extract relevant details in JSON format.",
29
+ )
30
+
31
+ # Launch the app
32
+ if __name__ == "__main__":
33
+ iface.launch()