first commit

app.py

@@ -11,13 +11,13 @@ if uploaded_file is not None:
         image = Image.open(uploaded_file)
         st.image(image, caption='Uploaded Image.', use_column_width=True)
         st.write("")
-        num_characters = st.text_input("Number of characters", "20")
+        num_tokens = st.text_input("Number of new tokens", "20")
         prompt = st.chat_input("Detect the main object in the image. The image is a")
-        num_characters = int(num_characters)
+        num_tokens = int(num_tokens)
         if prompt:
             st.write(f"User: {prompt}")
             with st.spinner('Processing...'):
-                label = kosmos.single_image_classification(image, prompt, num_characters)
+                label = kosmos.single_image_classification(image, prompt, num_tokens)
             st.write(f"Model: {label}")
     except Exception as e:
         st.error(f"An error occurred: {e}")

kosmos.py

@@ -21,7 +21,9 @@ def single_image_classification(image, prompt="", max_new_tokens=30):
     )
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    caption, entities = processor.post_process_generation(generated_text)
-    print("ROBOT:", caption)
-    print("ENTITIES: ", entities)
-    return caption
+    print("GENERATED:", generated_text)
+
+    processed_text = processor.post_process_generation(generated_text, cleanup_and_extract=False)
+
+    print("PROCESSED:", processed_text)
+    return processed_text