Hieucyber2208 commited on
Commit
8072d6c
·
verified ·
1 Parent(s): bb5f3c7

Update src/generation/llm.py

Browse files
Files changed (1) hide show
  1. src/generation/llm.py +1 -5
src/generation/llm.py CHANGED
@@ -99,12 +99,8 @@ class LLM:
99
  )
100
  # Decode the generated tokens
101
  response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
102
- # Strip any system/user metadata
103
- response = response.replace("system", "").replace("user", "").replace("assistant", "")
104
- # Remove any extra whitespace or unwanted tokens
105
- response = " ".join(response.split()).strip()
106
  print("Response generated successfully!")
107
- return response.strip()
108
  except Exception as e:
109
  raise RuntimeError(f"Failed to generate response: {str(e)}")
110
 
 
99
  )
100
  # Decode the generated tokens
101
  response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 
 
 
102
  print("Response generated successfully!")
103
+ return response
104
  except Exception as e:
105
  raise RuntimeError(f"Failed to generate response: {str(e)}")
106