embedding_model/data_preprocess/modify_question_model.py
2025-11-16 15:30:36 +00:00

104 lines
3.2 KiB
Python

from typing import List, Dict, Any
import json
import asyncio
import aiohttp
import time
import re
model_url = "http://192.168.130.206:4001/v1"
model = "google/gemma-3-27b-it"
class LLMModel:
def __init__(self):
self.instruction = """
You are a helpful assistant that help to me to modify and change the input question.
I will give you a question and its text and you must replace the words of question with synonyms or similar words.
## Important:
- replace the words of question with synonyms or similar words.
-the question must be in persian language.
return the question nothing else.
"""
async def run_llm(self, session, question, text):
"""
Run the llm model.
Args:
session: The session to use for the request.
question: The question to evaluate the text.
text: The text to evaluate.
Returns:
The result of the text.
"""
headers = {"Content-Type": "application/json"}
input_message = f"""{{"question": "{question}", "text": "{text}"}}"""
messages = [{"role": "system", "content": self.instruction}, {"role": "user", "content": input_message}]
payload = {
"model": model,
"messages": messages,
"max_tokens": 100
}
try:
async with session.post(model_url + "/chat/completions", headers=headers, json=payload) as resp:
resp.raise_for_status()
response = await resp.json()
result = response['choices'][0]['message']['content']
print(f"question: {question}")
print(f"result: {result}")
print("--------------------------------")
return result
except Exception as e:
try:
print(f"Error in llm model {response}: {e}")
except:
print(f"Error in llm model: {e}")
return ""
async def run_llm_async(self, question_list, text_list):
"""
Send all chunk requests concurrently.
Args:
question_list: The list of questions.
text_list: The list of texts.
Returns:
The list of results.
"""
async with aiohttp.ClientSession() as session:
tasks = [self.run_llm(session, question, text) for question, text in zip(question_list, text_list)]
results = await asyncio.gather(*tasks)
return results
def modify_question_llm(self, query_list: List[str], text_list: List[str]) -> List[Dict[str, Any]]:
"""
Modify question of the documents based on the query using the LLM model.
Args:
query_list: The list of queries.
text_list: The list of texts.
Returns:
The list of modified questions.
"""
if not text_list:
return []
start_time = time.time()
results = asyncio.run(self.run_llm_async(query_list, text_list))
end_time = time.time()
# print(f"Time taken for llm model: {end_time - start_time} seconds")
return results