Update app.py

app.py

@@ -11,25 +11,22 @@ MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
-
-
+# Controllo della disponibilità della GPU
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
-
+# Caricamento del modello e del tokenizer se la GPU è disponibile
 if torch.cuda.is_available():
     model_id = "AndreaAlessandrelli4/AvvoChat_AITA_v04"
-    model_vett_id = "intfloat/multilingual-e5-large"
     model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", load_in_4bit=True)
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     tokenizer.use_default_system_prompt = False
 
-
 @spaces.GPU
 def generate(
     message: str,
     chat_history: list[tuple[str, str]],
-    system_prompt: str,
+    system_prompt: str = "",
     max_new_tokens: int = 1024,
     temperature: float = 0.01,
     top_p: float = 0.9,
@@ -39,7 +36,7 @@ def generate(
 ) -> Iterator[str]:
     conversation = []
     conversation.append({"role": "system", "content": 
-                         '''Sei un an assistente AI di nome 'AvvoChat' specializzato nel rispondere a domande riguardanti la legge Italiana. 
+                         '''Sei un assistente AI di nome 'AvvoChat' specializzato nel rispondere a domande riguardanti la legge Italiana. 
                          Rispondi in lingua italiana in modo chiaro, semplice ed esaustivo alle domande che ti vengono fornite.
                          Le risposte devono essere sintetiche e chiare di massimo 500 token o anche più corte. 
                          Firmati alla fine di ogni risposta '-AvvoChat'.'''})
@@ -58,7 +55,7 @@ def generate(
         {"input_ids": input_ids},
         streamer=streamer,
         max_new_tokens=max_new_tokens,
-        do_sample=False,
+        do_sample=do_sample,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
@@ -73,69 +70,68 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
+image_path = "/home/a.alessandrelli/LLM_Dante/Raccolta_doc_index/AvvoVhat.png"
 
-chat_interface = gr.ChatInterface(
-    fn=generate,
-    additional_inputs=[
-        #gr.Textbox(label="System prompt", lines=6),
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
-        ),
-        gr.Slider(
-            label="Top-k",
-            minimum=1,
-            maximum=1000,
-            step=1,
-            value=50,
-        ),
-        gr.Checkbox(
-            label="Do-sample (False)",
-            value=False,
-        ),
-        gr.Slider(
-            label="Repetition penalty",
-            minimum=1.0,
-            maximum=2.0,
-            step=0.05,
-            value=1.2,
-        ),
-    ],
-    stop_btn=None,
-    examples=[
-        ["Posso fare un barbecue sul balcone di casa?"],
-        ["Posso essere multato se esco di casa senza documento d'identità?"],
-        ["Una persona single può adottare un bambino?"],
-        ["Posso usare un immagine creada con l'intelligenza artificiale?"],
-        ["Se il mio pallone da calcio cade in un giardino di un'abitazione privata, poss scavalcare il concello per riprendermelo?"],
-    ],
-)
-
-with gr.Blocks(css="style.css") as demo:
-    #gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
+with gr.Blocks() as demo:
     gr.Markdown("# AvvoChat")
     gr.Markdown("Fai una domanda riguardante la legge italiana all'AvvoChat e ricevi una spiegazione semplice al tuo dubbio.")
-    gr.Image("AvvoVhat.png", label="AvvoChat Logo", width=50, height=200)  # Dimensioni dell'immagine ridotte
-    chat_interface.render()
-    
 
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    with gr.Row():
+        msg = gr.Textbox(show_label=False, placeholder="Inserisci il tuo messaggio...")
+        btn = gr.Button("Invia")
+
+    with gr.Row():
+        with gr.Column(scale=0.5, min_width=100):
+            gr.Image(image_path, label="AvvoChat Logo", width=50, height=200)
+        with gr.Column(scale=6):
+            chatbox = gr.Chatbot(label="AvvoChat")
+
+    def on_send_message(message, chat_history, max_new_tokens, temperature, top_p, top_k, do_sample, repetition_penalty):
+        response_stream = generate(
+            message=message,
+            chat_history=chat_history,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            do_sample=do_sample,
+            repetition_penalty=repetition_penalty,
+        )
+        full_response = ""
+        for response in response_stream:
+            full_response += response
+        chat_history.append((message, full_response))
+        return chat_history, ""
+
+    # Collegare il pulsante e la textbox alla funzione on_send_message
+    btn.click(
+        on_send_message,
+        inputs=[
+            msg,
+            chatbox,
+            gr.State(value=DEFAULT_MAX_NEW_TOKENS),
+            gr.State(value=0.6),
+            gr.State(value=0.9),
+            gr.State(value=50),
+            gr.State(value=False),
+            gr.State(value=1.2)
+        ],
+        outputs=[chatbox, msg],
+    )
+    msg.submit(
+        on_send_message,
+        inputs=[
+            msg,
+            chatbox,
+            gr.State(value=DEFAULT_MAX_NEW_TOKENS),
+            gr.State(value=0.6),
+            gr.State(value=0.9),
+            gr.State(value=50),
+            gr.State(value=False),
+            gr.State(value=1.2)
+        ],
+        outputs=[chatbox, msg],
+    )
 
+if __name__ == "__main__":
+    demo.launch(share=True)