Update app.py

app.py

@@ -9,14 +9,16 @@ import spaces
 
 processor = AutoProcessor.from_pretrained("ucsahin/TraVisionLM-base", trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained("ucsahin/TraVisionLM-base", trust_remote_code=True)
+model_dpo = AutoModelForCausalLM.from_pretrained("ucsahin/TraVisionLM-DPO", trust_remote_code=True)
 
 model.to("cuda:0")
+model_dpo.to("cuda:0")
 
 @spaces.GPU
 def bot_streaming(message, history, max_tokens, temperature, top_p, top_k, repetition_penalty):
     print(max_tokens, temperature, top_p, top_k, repetition_penalty)
-    if message.files:
-        image = message.files[-1].path
+    if message['files']:
+        image = message['files'][-1]['path']
     else:
         # if there's no image uploaded for this turn, look for images in the past turns
         for hist in history:
@@ -26,11 +28,10 @@ def bot_streaming(message, history, max_tokens, temperature, top_p, top_k, repet
     if image is None:
         gr.Error("Lütfen önce bir resim yükleyin.")
 
-    prompt = f"{message.text}"
+    prompt = f"{message['text']}"
     image = Image.open(image).convert("RGB")
     inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda:0")
 
-    streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True})
     generation_kwargs = dict(
         inputs, streamer=streamer, max_new_tokens=max_tokens, 
         do_sample=True, temperature=temperature, top_p=top_p, 
@@ -38,18 +39,15 @@ def bot_streaming(message, history, max_tokens, temperature, top_p, top_k, repet
     )
     generated_text = ""
 
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+    model_outputs = model.generate(**generation_kwargs)
+    dpo_outputs = model_dpo.generate(**generation_kwargs)
 
-    text_prompt = f"{message.text}\n"
-
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        generated_text_without_prompt = buffer[len(text_prompt):]
+    model_output_text = processor.decode(model_outputs, skip_special_tokens=True)[len(prompt)+1]
+    dpo_output_text = processor.decode(dpo_outputs, skip_special_tokens=True)[len(prompt)+1]
     
-        time.sleep(0.04)
-        yield generated_text_without_prompt
+    generated_text = f"<h3>Base model cevabı:</h3>\n{model_output_text}\n<h3>DPO model cevabı:</h3>\n{dpo_output_text}"
+
+    return generated_text
 
 
 gr.set_static_paths(paths=["static/images/"])