| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline |
|
|
| |
| MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct" |
|
|
| |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True) |
| model = AutoModelForCausalLM.from_pretrained( |
| MODEL_NAME, |
| device_map="auto", |
| trust_remote_code=True |
| ) |
|
|
| |
| qwen_pipeline = pipeline( |
| "text-generation", |
| model=model, |
| tokenizer=tokenizer |
| ) |
|
|
| def generate_response(retrieved_texts, query, max_new_tokens=200): |
| context = "\n".join(retrieved_texts) |
| prompt = f"This is the detail about the image:\n{context}\n\nQuestion: {query}\nAnswer:" |
| |
| result = qwen_pipeline(prompt, max_new_tokens=max_new_tokens, ...) |
| generated_text = result[0]["generated_text"] |
| |
| if "Answer:" in generated_text: |
| final_answer = generated_text.split("Answer:")[-1].strip() |
| else: |
| final_answer = generated_text |
|
|
| return final_answer |
|
|
|
|