@@ -35,7 +35,7 @@ def get_inference_client(self):
3535
3636 def translate_code (self , source_code : str , source_lang : str , target_lang : str ) -> str :
3737 """
38- Translate code from one language to another using an instruct model.
38+ Translate code from one language to another using CodeLlama-34b- instruct
3939
4040 Args:
4141 source_code: Code to translate
@@ -48,41 +48,42 @@ def translate_code(self, source_code: str, source_lang: str, target_lang: str) -
4848 try :
4949 client = self .get_inference_client ()
5050
51+ # Create prompt for code translation
52+ prompt = f"""Translate the following { source_lang } code to { target_lang } .
53+ Only output the translated code without any explanations or markdown formatting.
54+
55+ { source_lang } code:
56+ ```
57+ { source_code }
58+ ```
59+
60+ { target_lang } code:
61+ ```"""
62+
5163 logger .info (f"Translating code from { source_lang } to { target_lang } " )
5264
53- response = client .chat .completions .create (
65+ # Use completions endpoint for CodeLlama
66+ response = client .completions .create (
5467 model = config .INFERENCE_MODEL_NAME ,
55- messages = [
56- {
57- "role" : "system" ,
58- "content" : (
59- f"You are an expert code translator. "
60- f"Translate { source_lang } code to { target_lang } . "
61- f"Output only the translated code with no explanations or markdown."
62- ),
63- },
64- {
65- "role" : "user" ,
66- "content" : f"Translate this { source_lang } code to { target_lang } :\n \n { source_code } " ,
67- },
68- ],
68+ prompt = prompt ,
6969 max_tokens = config .LLM_MAX_TOKENS ,
7070 temperature = config .LLM_TEMPERATURE ,
71+ stop = ["```" ] # Stop at closing code block
7172 )
7273
73- if response . choices :
74- translated_code = response .choices [ 0 ]. message . content . strip ()
75- # Strip markdown code fences if model wraps output anyway
76- if translated_code . startswith ( "```" ):
77- lines = translated_code . splitlines ()
78- translated_code = " \n " . join (
79- lines [ 1 : - 1 ] if lines [ - 1 ]. strip () == "```" else lines [ 1 :]
80- ). strip ()
81- logger .info (f"Successfully translated code ( { len ( translated_code ) } characters) " )
82- return translated_code
83-
84- logger .error (f"Empty choices in response: { response } " )
85- return ""
74+ # Handle response structure
75+ if hasattr ( response , 'choices' ) and len ( response .choices ) > 0 :
76+ choice = response . choices [ 0 ]
77+ if hasattr ( choice , 'text' ):
78+ translated_code = choice . text . strip ()
79+ logger . info ( f"Successfully translated code ( { len ( translated_code ) } characters)" )
80+ return translated_code
81+ else :
82+ logger .error (f"Unexpected response structure: { type ( choice ) } , { choice } " )
83+ return ""
84+ else :
85+ logger .error (f"Unexpected response: { type ( response ) } , { response } " )
86+ return ""
8687 except Exception as e :
8788 logger .error (f"Error translating code: { str (e )} " , exc_info = True )
8889 raise
0 commit comments