eagleswim commited on
Commit
c5369d3
·
verified ·
1 Parent(s): f424bed

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -58
app.py CHANGED
@@ -10,11 +10,7 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
10
  # Initialize vLLM with CPU configuration
11
  vllm_model = LLM(model=model_name, tensor_parallel_size=1, device="cpu")
12
 
13
- # Load the OCR model and processor
14
- ocr_model_name = "microsoft/trocr-small-handwritten"
15
- ocr_model = VisionEncoderDecoderModel.from_pretrained(ocr_model_name)
16
- ocr_processor = TrOCRProcessor.from_pretrained(ocr_model_name)
17
- #ocr_processor = AutoProcessor.from_pretrained(ocr_model_name)
18
 
19
  def generate_response(prompt, max_tokens, temperature, top_p):
20
  # Define sampling parameters
@@ -31,62 +27,11 @@ def generate_response(prompt, max_tokens, temperature, top_p):
31
  generated_text = output[0].outputs[0].text
32
  return generated_text
33
 
34
- def ocr_image(image_path):
35
- # Open the image from the file path
36
- image = Image.open(image_path).convert("RGB")
37
-
38
- # Preprocess the image for the OCR model
39
- pixel_values = ocr_processor(images=image, return_tensors="pt").pixel_values
40
-
41
- # Perform OCR inference
42
- outputs = ocr_model.generate(pixel_values)
43
-
44
- # Decode the generated tokens into text
45
- text = ocr_processor.batch_decode(outputs, skip_special_tokens=True)[0]
46
- return text
47
 
48
  # Gradio UI
49
  with gr.Blocks() as demo:
50
- gr.Markdown("# 🚀 Hugging Face Integration with vLLM and OCR (CPU)")
51
- gr.Markdown("Upload an image to extract text using OCR or generate text using the vLLM integration.")
52
-
53
- with gr.Tab("Text Generation"):
54
- with gr.Row():
55
- with gr.Column():
56
- prompt_input = gr.Textbox(
57
- label="Prompt",
58
- placeholder="Enter your prompt here...",
59
- lines=3,
60
- )
61
- max_tokens = gr.Slider(
62
- label="Max Tokens",
63
- minimum=10,
64
- maximum=500,
65
- value=100,
66
- step=10,
67
- )
68
- temperature = gr.Slider(
69
- label="Temperature",
70
- minimum=0.1,
71
- maximum=1.0,
72
- value=0.7,
73
- step=0.1,
74
- )
75
- top_p = gr.Slider(
76
- label="Top P",
77
- minimum=0.1,
78
- maximum=1.0,
79
- value=0.9,
80
- step=0.1,
81
- )
82
- submit_button = gr.Button("Generate")
83
-
84
- with gr.Column():
85
- output_text = gr.Textbox(
86
- label="Generated Text",
87
- lines=10,
88
- interactive=False,
89
- )
90
 
91
  submit_button.click(
92
  generate_response,
@@ -116,6 +61,11 @@ with gr.Blocks() as demo:
116
  inputs=[image_input],
117
  outputs=ocr_output,
118
  )
 
 
 
 
 
119
 
120
  # Launch the app
121
  demo.launch()
 
10
  # Initialize vLLM with CPU configuration
11
  vllm_model = LLM(model=model_name, tensor_parallel_size=1, device="cpu")
12
 
13
+
 
 
 
 
14
 
15
  def generate_response(prompt, max_tokens, temperature, top_p):
16
  # Define sampling parameters
 
27
  generated_text = output[0].outputs[0].text
28
  return generated_text
29
 
30
+
 
 
 
 
 
 
 
 
 
 
 
 
31
 
32
  # Gradio UI
33
  with gr.Blocks() as demo:
34
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
35
 
36
  submit_button.click(
37
  generate_response,
 
61
  inputs=[image_input],
62
  outputs=ocr_output,
63
  )
64
+ prompt =gr.Textbox()
65
+ max_tokens = gr.Textbox()
66
+ temperature = gr.Textbox()
67
+ top_p = gr.Textbox()
68
+ demo=gr.Interface(generate_response, inputs=[prompt, max_tokens,temperature, top_p], outputs="text")
69
 
70
  # Launch the app
71
  demo.launch()