Mori-kamiyama commited on
Commit
c44a239
·
1 Parent(s): 49df058

モデルを変えてみた

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -49,8 +49,12 @@ index = faiss.IndexFlatIP(dimension) # Cosine用にnormalize済ならこれ
49
  index.add(np.array(doc_embeddings))
50
 
51
  def generate_text(prompt):
52
- result = llm(search(prompt))
53
- return result['choices'][0]['text']
 
 
 
 
54
 
55
  def search(query):
56
  query_embedding = model.encode([query], normalize_embeddings=True)
 
49
  index.add(np.array(doc_embeddings))
50
 
51
  def generate_text(prompt):
52
+ full_prompt = search(prompt)
53
+ input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.cuda() # または .to("cuda") / .to("mps") / .to("cpu") に変更
54
+ with torch.no_grad():
55
+ output_ids = llm.generate(input_ids, max_new_tokens=256)
56
+ result_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
57
+ return result_text
58
 
59
  def search(query):
60
  query_embedding = model.encode([query], normalize_embeddings=True)