Update app.py

app.py

@@ -10,16 +10,17 @@ HF_TOKEN = os.getenv("HF_TOKEN")  # For Hugging Face Spaces, set this as a Secre
 
 # Global variable to store the pipeline
 text_generator_pipeline = None
-model_load_error = None
+model_load_error = None # To store any error message during model loading
 
 # --- Hugging Face Login and Model Loading ---
 def load_model_and_pipeline():
     global text_generator_pipeline, model_load_error
     if text_generator_pipeline is not None:
+        print("Model already loaded.")
         return True # Already loaded
 
     if not HF_TOKEN:
-        model_load_error = "Hugging Face token (HF_TOKEN) not found in Space secrets. Please add it."
+        model_load_error = "Hugging Face token (HF_TOKEN) not found in Space secrets. Please add it and restart the Space."
         print(f"ERROR: {model_load_error}")
         return False
 
@@ -29,30 +30,31 @@ def load_model_and_pipeline():
         print("Login successful.")
 
         print(f"Loading tokenizer for {MODEL_ID}...")
-        # trust_remote_code is necessary for some models that define custom architectures/code
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            trust_remote_code=True,
+            use_fast=False  # As recommended by the model card
+        )
+        # Llama models often don't have a pad token set by default
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
         print("Tokenizer loaded.")
 
         print(f"Loading model {MODEL_ID}...")
         # For large models, specify dtype and device_map
-        # device_map="auto" will try to use GPU if available, otherwise CPU
-        # torch_dtype="auto" or torch.bfloat16 (if supported by hardware) can save memory
-        # On CPU Spaces (free tier), this will be VERY slow or might OOM.
-        # You might need to use quantization (e.g., bitsandbytes) for CPU, but that's more complex.
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
             trust_remote_code=True,
-            torch_dtype="auto", # or torch.bfloat16 if on A10G or similar
-            device_map="auto" # "auto" is good for single/multi GPU or CPU fallback
+            torch_dtype=torch.bfloat16, # Use bfloat16 for better performance and memory if supported
+            device_map="auto"           # Automatically distribute model across available GPUs/CPU
         )
         print("Model loaded.")
 
-        # MIIA is an instruct/chat model, so text-generation is the appropriate task
         text_generator_pipeline = pipeline(
             "text-generation",
             model=model,
             tokenizer=tokenizer,
-            # device=0 if torch.cuda.is_available() else -1 # device_map handles this
+            # device_map="auto" handles device placement, so no need for device=0 here
         )
         print("Text generation pipeline created successfully.")
         model_load_error = None
@@ -71,21 +73,17 @@ def analyze_text(text_input, file_upload, custom_instruction, max_new_tokens, te
         if model_load_error:
             return f"Model not loaded. Error: {model_load_error}"
         else:
-            return "Model is not loaded. Please ensure HF_TOKEN is set and the Space has enough resources."
+            return "Model is not loaded or still loading. Please check Space logs for errors (especially OOM) and ensure HF_TOKEN is set and you've accepted model terms. If on CPU, it may take a very long time or fail due to memory."
 
     content_to_analyze = ""
     if file_upload is not None:
         try:
-            # file_upload is a TemporaryFileWrapper object, .name gives the path
             with open(file_upload.name, 'r', encoding='utf-8') as f:
                 content_to_analyze = f.read()
-            if not content_to_analyze.strip() and not text_input.strip(): # if file is empty and no text input
+            if not content_to_analyze.strip() and not text_input.strip():
                  return "Uploaded file is empty and no direct text input provided. Please provide some text."
-            elif not content_to_analyze.strip() and text_input.strip(): # if file empty but text input has content
+            elif not content_to_analyze.strip() and text_input.strip():
                 content_to_analyze = text_input
-            # If file has content, it will be used. If user also typed, file content takes precedence.
-            # We could add logic to concatenate or choose, but this is simpler.
-
         except Exception as e:
             return f"Error reading uploaded file: {str(e)}"
     elif text_input:
@@ -96,61 +94,61 @@ def analyze_text(text_input, file_upload, custom_instruction, max_new_tokens, te
     if not content_to_analyze.strip():
         return "Input text is empty."
 
-    # FastwebMIIA is an instruct model. It expects prompts like Alpaca.
-    # Structure:
-    # Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-    # ### Instruction:
-    # {your instruction}
-    # ### Input:
-    # {your text}
-    # ### Response:
-    # {model generates this}
-
-    prompt = f"""Di seguito è riportata un'istruzione che descrive un task, abbinata a un input che fornisce un contesto più ampio. Scrivi una risposta che completi la richiesta in modo appropriato.
-
-### Istruzione:
-{custom_instruction}
+    # Using Llama 2 Chat Format
+    # <s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{user_prompt} [/INST]
+    # For text analysis, the "instruction" is the user_prompt, and the "text_input" is part of it.
 
-### Input:
-{content_to_analyze}
+    system_prompt = "You are a helpful AI assistant specialized in text analysis. Perform the requested task on the provided text."
+    user_message = f"{custom_instruction}\n\nHere is the text:\n```\n{content_to_analyze}\n```"
 
-### Risposta:"""
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": user_message}
+    ]
 
-    # For English, you might change the preamble:
-    # prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-    # ### Instruction:
-    # {custom_instruction}
-    # ### Input:
-    # {content_to_analyze}
-    # ### Response:"""
+    try:
+        # Use tokenizer.apply_chat_template if available (transformers >= 4.34.0)
+        prompt = text_generator_pipeline.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    except Exception as e:
+        print(f"Warning: Could not use apply_chat_template ({e}). Falling back to manual formatting.")
+        # Manual Llama 2 chat format
+        prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{user_message} [/INST]"
 
 
     print(f"\n--- Sending to Model ---")
-    print(f"Prompt:\n{prompt}")
+    print(f"Full Prompt:\n{prompt}")
     print(f"Max New Tokens: {max_new_tokens}, Temperature: {temperature}, Top P: {top_p}")
     print("------------------------\n")
 
     try:
-        # Note: text-generation pipelines often return the prompt + completion.
-        # We might need to strip the prompt from the output if desired.
         generated_outputs = text_generator_pipeline(
             prompt,
             max_new_tokens=int(max_new_tokens),
             do_sample=True,
-            temperature=float(temperature) if float(temperature) > 0 else 0.7, # temp 0 means greedy
+            temperature=float(temperature) if float(temperature) > 0.01 else 0.01, # Temperature 0 can be problematic
             top_p=float(top_p),
-            num_return_sequences=1
+            num_return_sequences=1,
+            eos_token_id=text_generator_pipeline.tokenizer.eos_token_id,
+            pad_token_id=text_generator_pipeline.tokenizer.pad_token_id # Use the set pad_token
         )
-        response = generated_outputs[0]['generated_text']
+        response_full = generated_outputs[0]['generated_text']
 
-        # Often, the response includes the prompt. Let's try to return only the new part.
-        # The model should generate text after "### Risposta:"
-        answer_marker = "### Risposta:"
-        if answer_marker in response:
-            return response.split(answer_marker, 1)[1].strip()
+        # Extract only the assistant's response part
+        # The model's actual answer starts after the [/INST] token.
+        answer_marker = "[/INST]"
+        if answer_marker in response_full:
+            response_text = response_full.split(answer_marker, 1)[1].strip()
         else:
-            # Fallback if the marker isn't found (shouldn't happen with good prompting)
-            return response # Or you could try to remove the original prompt string
+            # Fallback if the full prompt wasn't returned, might happen with some pipeline configs
+            # or if the model didn't fully adhere to the template in its output.
+            # This is less ideal, but better than nothing.
+            response_text = response_full.replace(prompt, "").strip() # Try to remove the input prompt
+
+        return response_text
 
     except Exception as e:
         return f"Error during text generation: {str(e)}"
@@ -162,31 +160,35 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     Test the capabilities of the `{MODEL_ID}` model for text analysis tasks on Italian or English texts.
     Provide an instruction and your text (directly or via upload).
     **Important:** Model loading can take a few minutes, especially on the first run or on CPU.
-    This app is best run on a Hugging Face Space with GPU resources for this model size.
+    This app is best run on a Hugging Face Space with GPU resources (e.g., T4-small or A10G-small) for this 7B model.
     """)
 
     with gr.Row():
-        status_textbox = gr.Textbox(label="Model Status", value="Attempting to load model...", interactive=False)
+        status_textbox = gr.Textbox(label="Model Status", value="Initializing...", interactive=False, scale=3)
+        current_hardware = os.getenv("SPACE_HARDWARE", "Unknown (likely local or unspecified)")
+        gr.Markdown(f"Running on: **{current_hardware}**")
+
 
     with gr.Tab("Text Input & Analysis"):
         with gr.Row():
             with gr.Column(scale=2):
                 instruction_prompt = gr.Textbox(
-                    label="Instruction for the Model (e.g., 'Riassumi questo testo', 'Identify main topics', 'Translate to English')",
+                    label="Instruction for the Model (Cosa vuoi fare con il testo?)",
                     value="Riassumi questo testo in 3 frasi concise.",
-                    lines=3
+                    lines=3,
+                    placeholder="Example: Riassumi questo testo. / Summarize this text. / Estrai le entità nominate. / Identify named entities."
                 )
-                text_area_input = gr.Textbox(label="Enter Text Directly", lines=10, placeholder="Paste your text here...")
-                file_input = gr.File(label="Or Upload a Document (.txt)", file_types=['.txt'])
+                text_area_input = gr.Textbox(label="Enter Text Directly / Inserisci il testo direttamente", lines=10, placeholder="Paste your text here or upload a file below...")
+                file_input = gr.File(label="Or Upload a Document (.txt) / O carica un documento (.txt)", file_types=['.txt'])
             with gr.Column(scale=3):
-                output_text = gr.Textbox(label="Model Output", lines=20, interactive=False)
+                output_text = gr.Textbox(label="Model Output / Risultato del Modello", lines=20, interactive=False)
 
         with gr.Accordion("Advanced Generation Parameters", open=False):
-            max_new_tokens_slider = gr.Slider(minimum=50, maximum=1024, value=256, step=10, label="Max New Tokens")
-            temperature_slider = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature (higher is more creative)")
+            max_new_tokens_slider = gr.Slider(minimum=10, maximum=2048, value=256, step=10, label="Max New Tokens")
+            temperature_slider = gr.Slider(minimum=0.01, maximum=2.0, value=0.7, step=0.01, label="Temperature (higher is more creative, 0.01 for more deterministic)")
             top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top P (nucleus sampling)")
 
-        analyze_button = gr.Button("🧠 Analyze Text", variant="primary")
+        analyze_button = gr.Button("🧠 Analyze Text / Analizza Testo", variant="primary")
 
     analyze_button.click(
         fn=analyze_text,
@@ -197,10 +199,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     # Load the model when the app starts.
     # This will update the status_textbox after attempting to load.
     def startup_load_model():
+        print("Gradio app starting, attempting to load model...")
         if load_model_and_pipeline():
             return "Model loaded successfully and ready."
         else:
-            return f"Failed to load model. Error: {model_load_error or 'Unknown error during startup.'}"
+            return f"Failed to load model. Error: {model_load_error or 'Unknown error during startup. Check Space logs.'}"
 
     demo.load(startup_load_model, outputs=status_textbox)
 
@@ -208,18 +211,21 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 if __name__ == "__main__":
     # For local testing (ensure HF_TOKEN is set as an environment variable or you're logged in via CLI)
     # You would run: HF_TOKEN="your_hf_token_here" python app.py
-    # If not set, it will fail unless you've done `huggingface-cli login`
     if not HF_TOKEN and "HF_TOKEN" not in os.environ:
         print("WARNING: HF_TOKEN environment variable not set.")
         print("For local execution, either set HF_TOKEN or ensure you are logged in via 'huggingface-cli login'.")
-        # Attempt to use CLI login if available
         try:
-            HF_TOKEN = huggingface_hub.HfApi().token
-            if HF_TOKEN:
+            from huggingface_hub import HfApi
+            hf_api = HfApi()
+            token = hf_api.token
+            if token:
+                os.environ['HF_TOKEN'] = token # Set it for the current process
+                HF_TOKEN = token # also update the global variable used by the script
                 print("Using token from huggingface-cli login.")
             else:
                 print("Could not retrieve token from CLI login. Model access might fail.")
         except Exception as e:
             print(f"Could not check CLI login status: {e}. Model access might fail.")
 
+    print("Launching Gradio interface...")
     demo.queue().launch(debug=True, share=False) # share=True for public link if local