Spaces:
Running
Running
Mr-Vicky-01
commited on
Update llm.py
Browse files
llm.py
CHANGED
@@ -1,19 +1,17 @@
|
|
1 |
-
import google.generativeai as genai
|
2 |
-
from
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
img = Image.open(image)
|
18 |
-
response = self.model.generate_content([prompt, img])
|
19 |
return response.text
|
|
|
1 |
+
import google.generativeai as genai
|
2 |
+
from PIL import Image
|
3 |
+
import os
|
4 |
+
|
5 |
+
os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
|
6 |
+
genai.configure(api_key=os.environ["GOOGLE_API_KEY"])
|
7 |
+
|
8 |
+
class Model:
|
9 |
+
def __init__(self) -> None:
|
10 |
+
self.model = genai.GenerativeModel('gemini-1.5-flash-latest')
|
11 |
+
|
12 |
+
def get_response(self, image):
|
13 |
+
prompt = """You are an intelligent document creator. Could you please extract the words from the given screenshot and provide me document text that matches exact screenshot font and look
|
14 |
+
important note: if the screenshot not contain any text means you must say 'please upload a valid screenshot.'"""
|
15 |
+
img = Image.open(image)
|
16 |
+
response = self.model.generate_content([prompt, img])
|
|
|
|
|
17 |
return response.text
|