vithacocf commited on
Commit
39d84f4
·
verified ·
1 Parent(s): 19ac777

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -60
app.py CHANGED
@@ -1,69 +1,52 @@
1
  import gradio as gr
2
- # from transformers import AutoProcessor, AutoModelForVision2Seq
3
  from PIL import Image
4
- # import torch
5
-
6
- # device = "cuda" if torch.cuda.is_available() else "cpu"
7
- # torch.cuda.empty_cache()
8
-
9
- # model_id = "prithivMLmods/Camel-Doc-OCR-062825"
10
- # processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
11
- # model = AutoModelForVision2Seq.from_pretrained(
12
- # model_id,
13
- # torch_dtype=torch.float16 if device == "cuda" else torch.float32,
14
- # trust_remote_code=True
15
- # ).to(device)
16
-
17
- # def predict(image, prompt=None):
18
- # image = image.convert("RGB")
19
-
20
- # # Cực kỳ quan trọng: text="" bắt buộc phải có
21
- # inputs = processor(images=image, text="", return_tensors="pt").to(device)
22
- # # In debug để kiểm tra input_ids
23
- # print(">>> input_ids shape:", inputs.input_ids.shape)
24
- # generated_ids = model.generate(
25
- # **inputs,
26
- # max_new_tokens=512,
27
- # do_sample=False,
28
- # use_cache=False, # ✅ Thêm dòng này để fix lỗi cache_position
29
- # eos_token_id=processor.tokenizer.eos_token_id,
30
- # pad_token_id=processor.tokenizer.pad_token_id
31
- # )
32
-
33
- # result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
34
- # return result
35
-
36
- # Hàm mock xử lý ảnh — chỉ để test UI
37
- # def mock_predict(image, prompt=None):
38
- # return f"Fake OCR result for image. Prompt: {prompt or 'N/A'}"
39
-
40
- # demo = gr.Interface(
41
- # # fn=predict,
42
- # fn=mock_predict,
43
- # inputs=[
44
- # gr.Image(type="pil", label="Tải ảnh tài liệu lên"),
45
- # gr.Textbox(label="Gợi ý (tuỳ chọn)", placeholder="VD: Trích số hóa đơn")
46
- # ],
47
- # outputs="text",
48
- # title="Camel-Doc OCR - Trích xuất văn bản từ ảnh"
49
- # )
50
-
51
- # # if __name__ == "__main__":
52
- # # demo.launch()
53
- # demo.launch(share=False)
54
-
55
- def predict(image: Image.Image, prompt: str = "") -> str:
56
- return f"✅ UI OK! Prompt: {prompt}"
57
 
58
  demo = gr.Interface(
59
- fn=predict,
 
60
  inputs=[
61
- gr.Image(type="pil", label="Ảnh"),
62
- gr.Textbox(label="Prompt")
63
  ],
64
  outputs="text",
65
- title="Test Gradio UI"
66
  )
67
 
68
- # **BẮT BUỘC** gọi launch() ở cuối
69
- demo.launch(debug=True)
 
1
  import gradio as gr
2
+ from transformers import AutoProcessor, AutoModelForVision2Seq
3
  from PIL import Image
4
+ import torch
5
+
6
+ device = "cuda" if torch.cuda.is_available() else "cpu"
7
+ torch.cuda.empty_cache()
8
+
9
+ model_id = "prithivMLmods/Camel-Doc-OCR-062825"
10
+ processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
11
+ model = AutoModelForVision2Seq.from_pretrained(
12
+ model_id,
13
+ torch_dtype=torch.float16 if device == "cuda" else torch.float32,
14
+ trust_remote_code=True
15
+ ).to(device)
16
+
17
+ def predict(image, prompt=None):
18
+ image = image.convert("RGB")
19
+
20
+ # Cực kỳ quan trọng: text="" bắt buộc phải có
21
+ inputs = processor(images=image, text="", return_tensors="pt").to(device)
22
+ # In debug để kiểm tra input_ids
23
+ print(">>> input_ids shape:", inputs.input_ids.shape)
24
+ generated_ids = model.generate(
25
+ **inputs,
26
+ max_new_tokens=512,
27
+ do_sample=False,
28
+ use_cache=False, # ✅ Thêm dòng này để fix lỗi cache_position
29
+ eos_token_id=processor.tokenizer.eos_token_id,
30
+ pad_token_id=processor.tokenizer.pad_token_id
31
+ )
32
+
33
+ result = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
34
+ return result
35
+
36
+ Hàm mock xử lý ảnh — chỉ để test UI
37
+ def mock_predict(image, prompt=None):
38
+ return f"Fake OCR result for image. Prompt: {prompt or 'N/A'}"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
 
40
  demo = gr.Interface(
41
+ # fn=predict,
42
+ fn=mock_predict,
43
  inputs=[
44
+ gr.Image(type="pil", label="Tải ảnh tài liệu lên"),
45
+ gr.Textbox(label="Gợi ý (tuỳ chọn)", placeholder="VD: Trích số hóa đơn")
46
  ],
47
  outputs="text",
48
+ title="Camel-Doc OCR - Trích xuất văn bản từ ảnh"
49
  )
50
 
51
+ # if __name__ == "__main__":
52
+ # demo.launch()