extract_relevant_text

app.py

@@ -2,6 +2,7 @@ import os
 from flask import Flask, jsonify, request
 from flask_cors import CORS 
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig  
+import re
 
 # Set the HF_HOME environment variable to a writable directory
 os.environ["HF_HOME"] = "/workspace/huggingface_cache"  # Change this to a writable path in your space
@@ -20,7 +21,7 @@ def get_model_and_tokenizer(model_id):
     try:
         print(f"Loading tokenizer for model_id: {model_id}")
         # Load the tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
         tokenizer.pad_token = tokenizer.eos_token
         
         print(f"Loading model and for model_id: {model_id}")
@@ -31,6 +32,25 @@ def get_model_and_tokenizer(model_id):
     except Exception as e:
         print(f"Error loading model: {e}")
 
+def extract_relevant_text(response):
+    """
+    This function extracts the first 'user' and 'assistant' blocks between
+    <|im_start|> and <|im_end|> in the generated response.
+    """
+    # Regex to match content between <|im_start|> and <|im_end|> tags
+    pattern = re.compile(r"<\|im_start\|>(.*?)<\|im_end\|>", re.DOTALL)
+    matches = pattern.findall(response)
+
+    if len(matches) < 2:
+        return "Unable to extract sufficient data from the response."
+
+    # Assuming the first match is user and the second match is assistant
+    user_message = matches[0].strip()  # First <|im_start|> block
+    assistant_message = matches[1].strip()  # Second <|im_start|> block
+    
+    # Format the extracted result
+    return f"user: {user_message}\nassistant: {assistant_message}"
+
 def generate_response(user_input, model_id):
     prompt = formatted_prompt(user_input)
     
@@ -56,7 +76,7 @@ def generate_response(user_input, model_id):
         # Generate response
         outputs = model.generate(**inputs, generation_config=generation_config)
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response
+        return extract_relevant_text(response)
     except Exception as e:
         print(f"Error generating response: {e}")
         return "Error generating response."