Update app.py

app.py

@@ -1,11 +1,10 @@
 import os
 from threading import Thread
 from typing import Iterator
+
 import gradio as gr
 import spaces
 import torch
-import weaviate
-from sentence_transformers import SentenceTransformer
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 MAX_MAX_NEW_TOKENS = 2048
@@ -33,7 +32,6 @@ if torch.cuda.is_available():
     )
 
 
-
 def prompt_template(materiali, query):
     mat = ''
     for i, doc in enumerate(materiali):
@@ -67,13 +65,14 @@ def richiamo_materiali(query, vett_query, alpha=1.0, N_items=5):
 
 
 
+
 @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.1,
+    temperature: float = 0.6,
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
@@ -96,14 +95,15 @@ def generate(
     input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Chat troppo lunga superati {MAX_INPUT_TOKEN_LENGTH} tokens.")
+        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     input_ids = input_ids.to(model.device)
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+    streamer = TextIteratorStreamer(tokenizer, timeout=30.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         {"input_ids": input_ids},
         streamer=streamer,
         max_new_tokens=max_new_tokens,
+        do_sample=True,
         top_p=top_p,
         top_k=top_k,
         temperature=temperature,
@@ -123,14 +123,50 @@ chat_interface = gr.ChatInterface(
     fn=generate,
     chatbot=gr.Chatbot(height=400, label = "AvvoChat", show_copy_button=True, avatar_images=("users.jpg","AvvoVhat.png"), show_share_button=True),
     textbox=gr.Textbox(placeholder="Inserisci la tua domanda", container=False, scale=7),
+    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.Slider(
+            label="Repetition penalty",
+            minimum=1.0,
+            maximum=2.0,
+            step=0.05,
+            value=1.2,
+        ),
+    ],
     submit_btn ="Chiedi all'AvvoChat ",
     retry_btn = "Rigenera",
     undo_btn = None,
     clear_btn = "Pulisci chat",
     fill_height = True,
     theme = "gstaff/sketch",
-    #title="Avvo-Chat",
-    #description="""Fai una domanda riguardante la legge italiana all'AvvoChat e ricevi una spiegazione semplice al tuo dubbio.""",
     examples=[
         ["Posso fare una grigliata sul balcone di casa?"],
         ["Se esco di casa senza documento di identità posso essere multato?"],
@@ -151,4 +187,4 @@ with gr.Blocks(css="style.css") as demo:
 
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20).launch()