first commit

app.py

@@ -15,7 +15,7 @@ if uploaded_file is not None:
         if st.button('Classify'):
             st.write(f"User: {prompt}")
             with st.spinner('Processing...'):
-                label = kosmos.single_image_classification(image, prompt, 50)
+                label = kosmos.single_image_classification(image, prompt, 150)
             st.write(f"Model: {label}")
     except Exception as e:
         st.error(f"An error occurred: {e}")

kosmos.py

@@ -19,8 +19,8 @@ def single_image_classification(image, prompt="", max_new_tokens=30):
         use_cache=True,
         max_new_tokens=max_new_tokens,
     )
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    processed_text = processor.post_process_generation(generated_text, cleanup_and_extract=False)
-    print("ROBOT:", processed_text)
+    caption, entities = processor.post_process_generation(generated_text)
+    print("ROBOT:", caption)
     return processed_text