Update app.py

local_generate

app.py

@@ -30,6 +30,29 @@ except Exception as e:
   global_error =  f"Failed to load model: {str(e)}"
 
 @spaces.GPU()
+def local_generate(
+    prompt,
+    generation_config,
+    max_new_tokens,
+    do_sample,
+    top_p,
+    repetition_penalty,
+    temperature=temperature
+):
+    response_output = generator(
+        prompt,
+        generation_config=generation_config,
+        max_new_tokens=max_new_tokens,
+        do_sample=do_sample,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        temperature=temperature
+    )        
+    generated_text = response_output[0]['generated_text']
+    # Extract the assistant's response
+    result = generated_text[len(prompt):]    
+    return result
+
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -82,18 +105,15 @@ def respond(
         full_result = ''
         while ( (tokens_cnt < max_tokens) and (last_token_len > 0) ):
             # Generate the response
-            response_output = generator(
+            result = local_generate(
                 prompt,
                 generation_config=generator_conf,
-                max_new_tokens=tokens_inc,
+                max_new_tokens=max_tokens,
                 do_sample=True,
                 top_p=top_p,
                 repetition_penalty=1.2,
                 temperature=temperature
-            )        
-            generated_text = response_output[0]['generated_text']
-            # Extract the assistant's response
-            result = generated_text[len(prompt):]
+            )
             full_result = full_result + result
             prompt = prompt + result
             tokens_cnt = tokens_cnt + tokens_inc