Refactor app.py to implement streaming response generation for email creation. Removed tokenizer and model initialization in favor of using the InferenceClient for improved performance and reduced complexity. Updated conversation_predict function to yield streaming output, enhancing user experience with real-time email generation feedback. Added a new update_ui function to manage email display in the Streamlit interface, ensuring a smoother interaction flow.

app.py

@@ -78,15 +78,6 @@ def extract_cv_text(file):
 # Replace 'your_huggingface_token' with your actual Hugging Face access token
 access_token = os.getenv('API_KEY')
 
-# Initialize the tokenizer and model with the Hugging Face access token
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it", use_auth_token=access_token)
-model = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-2b-it",
-    torch_dtype=torch.bfloat16,
-    use_auth_token=access_token
-)
-model.eval()  # Set the model to evaluation mode
-
 # Initialize the inference client (if needed for other API-based tasks)
 client = InferenceClient(token=access_token)
 
@@ -118,23 +109,22 @@ Keep the tone professional, confident, and enthusiastic. Be concise but impactfu
 Email:"""
 
 def conversation_predict(input_text: str, cv_sections: Dict[str, str]):
-    """Generate a response using the model with improved prompting."""
+    """Generate a response using the model with streaming output."""
     prompt = create_email_prompt(input_text, cv_sections)
     
-    # Tokenize the input text
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-
-    # Generate a response with the model
-    outputs = model.generate(
-        input_ids,
-        max_new_tokens=2048,
+    # Use the streaming API
+    for response in client.text_generation(
+        model="google/gemma-2b-it",
+        prompt=prompt,
+        max_new_tokens=512,
         temperature=0.7,
         top_p=0.95,
-        do_sample=True
-    )
-
-    # Decode and return the generated response
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+        stream=True
+    ):
+        if hasattr(response, 'token'):  # Handle different response formats
+            yield response.token.text
+        else:
+            yield response.generated_text
 
 def respond(
     message: str,
@@ -194,18 +184,34 @@ with tab1:
         if isinstance(cv_sections, dict):
             st.success("CV uploaded and parsed successfully!")
         else:
-            st.error(cv_sections)  # Show error message if parsing failed
+            st.error(cv_sections)
     
     # Job description input
     st.markdown("### Job Description")
     message = st.text_area("Paste the job description here:", height=200)
     
+    # Call the updated UI function
+    update_ui()
+
+def update_ui():
+    # Create placeholder for the generated email
+    email_placeholder = st.empty()
+    
     # Generate button
     if st.button("Generate Email"):
         if message and cv_file and isinstance(cv_sections, dict):
-            response = conversation_predict(message, cv_sections)
-            st.markdown("### Generated Email:")
-            st.markdown(response)
+            email_text = ""
+            # Stream the response
+            for chunk in conversation_predict(message, cv_sections):
+                if chunk:
+                    email_text += chunk
+                    # Update the text area with each chunk
+                    email_placeholder.text_area(
+                        "Generated Email",
+                        value=email_text,
+                        height=400,
+                        key="email_output"
+                    )
         else:
             st.warning("Please upload a CV and enter a job description.")