demo.py 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. import numpy as np
  2. from models.model_loader import model_loader
  3. class AIService:
  4. # 创建AI处理类
  5. def __init__(self):
  6. self.text_model = model_loader.load_text_model()
  7. # 常用AI方法
  8. def text_to_embedding(self, text):
  9. # 1.文本转向量
  10. if isinstance(text, str):
  11. text = [text]
  12. embeddings = self.text_model.encode(text)
  13. # 返回numpy数组,不要转成list
  14. return embeddings
  15. def calculate_similarity(self, text1, text2):
  16. # 2.计算文本相似度
  17. emb1 = self.text_to_embedding(text1)
  18. emb2 = self.text_to_embedding(text2)
  19. # 确保是1D数组
  20. if len(emb1.shape) > 1:
  21. emb1 = emb1[0]
  22. if len(emb2.shape) > 1:
  23. emb2 = emb2[0]
  24. similarity = np.dot(emb1, emb2) / (
  25. np.linalg.norm(emb1) * np.linalg.norm(emb2)
  26. )
  27. return float(similarity)
  28. def batch_text_processing(self, texts):
  29. # 3.批量文本处理
  30. embeddings = self.text_to_embedding(texts)
  31. # 简单的文本分析(可扩展)
  32. results = []
  33. for i, text in enumerate(texts):
  34. results.append({
  35. 'text': text,
  36. 'embedding_size': embeddings.shape[1], # 使用shape获取维度
  37. 'text_length': len(text),
  38. 'embedding': embeddings[i].tolist() # 只在返回给前端时转成list
  39. })
  40. return results
  41. # 全局AI服务实例
  42. ai_service = AIService()